From df0d6ad62a8d82f69fae18f8449f63f8996d16cc Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Mon, 19 Oct 2026 18:15:24 +0000 Subject: [PATCH 01/11] Route ug copilot through a refreshing loopback auth proxy Copilot CLI holds its provider credential for the whole session, so the Databricks token ug passed at launch expired mid-session (the 30-minute refresher only rewrote ~/.copilot/ucode.env, which Copilot never reads). The token was also persisted in that file and its backup, and `ug revert` restored it. `ug copilot` now starts a per-session proxy on 127.0.0.1. Copilot gets the proxy URL and a random session secret; the proxy accepts the secret only on gateway model routes (/ai-gateway/{anthropic,mlflow/v1,openai/v1}/), swaps it for a fresh Databricks token (CLI profile or custom OAuth), and adds the ucode/copilot User-Agent, request tags and managed http_headers. A rejected token is refreshed and the request retried; when refresh itself fails the proxy answers with a "Databricks sign-in expired" error in the request's dialect and logs to ~/.ucode/logs/copilot-proxy.log instead of the TUI. - Configure and launch write no credential; earlier tokens are scrubbed from ucode.env, a token-bearing backup is dropped (also before revert), and ug no longer backs up its own file on reconfigure. - The refresher thread is gone, so a running session no longer resets current_workspace or rewrites state. - Ctrl-C no longer re-signals Copilot (send_signal(SIGINT) raised on Windows); ug waits for Copilot to exit and then stops the proxy. - Inherited COPILOT_PROVIDER_API_KEY(_COMMAND) / _TRANSPORT are dropped so they can't bypass the proxy. Tests: hermetic socket tests for the proxy (secret, path allowlist incl. encoded traversal, header swap, refresh-and-retry, actionable 401), Copilot unit tests, the Copilot e2e now runs through the proxy, and a new non-blocking Copilot integration lane (Linux + Windows) with Claude and GPT headless journeys that also assert no credential on disk. Co-authored-by: Isaac --- .github/workflows/integration.yml | 28 ++ README.md | 7 + scripts/run_integration.py | 36 +- src/ucode/agents/__init__.py | 2 +- src/ucode/agents/copilot.py | 212 +++++++---- src/ucode/cli.py | 5 +- src/ucode/gateway_proxy.py | 168 ++++++++- tests/AGENTS.md | 3 +- tests/README.md | 1 + tests/integration/README.md | 12 +- tests/integration/conftest.py | 3 +- tests/integration/pytest.ini | 1 + tests/integration/test_ug_copilot_headless.py | 99 +++++ tests/integration/utils/evidence.py | 9 + tests/test_agent_copilot.py | 357 ++++++++++++++---- tests/test_e2e.py | 56 ++- tests/test_gateway_proxy_integration.py | 187 ++++++++- tests/test_integration_contract.py | 3 +- tests/test_integration_runner.py | 20 + 19 files changed, 1023 insertions(+), 186 deletions(-) create mode 100644 tests/integration/test_ug_copilot_headless.py diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index d7ff07e4d..70aaff861 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -30,6 +30,10 @@ on: description: Exact OpenCode version (non-blocking lane) default: 1.18.31 required: true + copilot_version: + description: Exact GitHub Copilot CLI version (non-blocking lane) + default: 1.0.91 + required: true dependency: description: Optional exact dependency for reproduction (e.g. tomlkit==0.14.0) default: '' @@ -53,6 +57,7 @@ env: CLAUDE_VERSION: ${{ inputs.claude_version || '2.1.290' }} CODEX_VERSION: ${{ inputs.codex_version || '0.154.0' }} OPENCODE_VERSION: ${{ inputs.opencode_version || '1.18.31' }} + COPILOT_VERSION: ${{ inputs.copilot_version || '1.0.91' }} PACKAGE_INDEX: ${{ inputs.default_index || 'https://pypi.org/simple' }} jobs: @@ -211,6 +216,7 @@ jobs: claude) args=(--claude-version "$CLAUDE_VERSION") ;; codex) args=(--codex-version "$CODEX_VERSION") ;; opencode) args=(--opencode-version "$OPENCODE_VERSION") ;; + copilot) args=(--copilot-version "$COPILOT_VERSION") ;; esac if [[ -n "$DEPENDENCY" ]]; then args+=(--dependency "$DEPENDENCY") @@ -351,6 +357,28 @@ jobs: ARTIFACT_NAME: integration-full-opencode-${{ matrix.os }} steps: *live-steps + # Non-blocking pending stabilization, like OpenCode: continue-on-error and absent from the + # cujs `needs`. + copilot: + name: Full journeys · Copilot · ${{ matrix.os == 'windows-latest' && 'Windows' || 'Linux' }} (non-blocking) + if: ${{ !cancelled() && needs.workspace.result == 'success' }} + needs: workspace + continue-on-error: true + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04, windows-latest] + timeout-minutes: 30 + env: + UCODE_TEST_WORKSPACE: ${{ secrets.UCODE_TEST_WORKSPACE }} + DATABRICKS_BEARER: ${{ secrets.DATABRICKS_BEARER }} + DEPENDENCY: ${{ inputs.dependency }} + AGENT: copilot + TEST_MARKER: live and copilot + ARTIFACT_NAME: integration-full-copilot-${{ matrix.os }} + steps: *live-steps + # Non-blocking pending stabilization: continue-on-error and absent from the cujs `needs`. # Runs on the managed e2e workspace with an injected config that enables Claude/Codex but not # OpenCode (a service principal authenticates it, like the `managed` lane). diff --git a/README.md b/README.md index 3de696aa7..fe7f7b0a0 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,13 @@ because it overrides the model sent to the gateway. Restart Copilot through `ug` to change the wire model or API; in-session model selection does not rebuild its provider configuration. +`ug copilot` runs Copilot against a loopback proxy for the session. Copilot holds only +a random session secret; the proxy adds a fresh Databricks token to each gateway +request, so OAuth sessions keep working past token expiry without a restart, and no +token is written to `~/.copilot/ucode.env`. If the Databricks sign-in itself expires, +Copilot shows a "Databricks sign-in expired" error; run `databricks auth login` and +restart `ug copilot`. Proxy errors are logged to `~/.ucode/logs/copilot-proxy.log`. + Without a managed workspace config, `ug claude` automatically discovers gateway models for Claude Code's `/model` picker. Discovery defaults to `system.ai` when no provider or model location is selected. Use `--provider` or `--model-location` diff --git a/scripts/run_integration.py b/scripts/run_integration.py index 0124bfd0c..f252dd3f3 100644 --- a/scripts/run_integration.py +++ b/scripts/run_integration.py @@ -30,6 +30,7 @@ "claude": "@anthropic-ai/claude-code", "codex": "@openai/codex", "opencode": "opencode-ai", + "copilot": "@github/copilot", } WINDOWS_PATHEXT = ".COM;.EXE;.BAT;.CMD" # Claude exports no spans on Windows, where ug writes no machine-wide Claude settings. @@ -44,9 +45,18 @@ "claude": "test_ug_claude_headless.py::test_ug_claude_headless_prompt_argument", "codex": "test_ug_codex_headless.py::test_ug_codex_headless_prompt_argument", "opencode": "test_ug_opencode_headless.py::test_ug_opencode_headless_prompt_argument", + "copilot": "test_ug_copilot_headless.py::test_ug_copilot_headless_completes_file_task", } +def agent_version_output(agent: str, output: str) -> str: + """The agent's ``--version`` text, normalized for exact-version matching.""" + if agent == "copilot": + # "GitHub Copilot CLI 1.0.91." plus an update hint on later lines. + return output.splitlines()[0].removesuffix(".") if output else output + return output + + def redact_secrets(value: str, secrets: Iterable[str]) -> str: for secret in sorted({secret for secret in secrets if secret}, key=len, reverse=True): value = value.replace(secret, "") @@ -253,11 +263,13 @@ def arguments( parser.add_argument("--claude-version", type=exact_npm_version) parser.add_argument("--codex-version", type=exact_npm_version) parser.add_argument("--opencode-version", type=exact_npm_version) + parser.add_argument("--copilot-version", type=exact_npm_version) parser.add_argument("--claude-model", default=environment.get("UG_INTEGRATION_CLAUDE_MODEL")) parser.add_argument("--codex-model", default=environment.get("UG_INTEGRATION_CODEX_MODEL")) parser.add_argument( "--opencode-model", default=environment.get("UG_INTEGRATION_OPENCODE_MODEL") ) + parser.add_argument("--copilot-model", default=environment.get("UG_INTEGRATION_COPILOT_MODEL")) parser.add_argument( "--claude-provider", default="main.ucode.ci_e2e_anthropic_nonrelay_mps", @@ -356,9 +368,12 @@ def arguments( args.pytest_args.extend([flag, str(value)]) if selected.x: args.pytest_args.append("-x") - if not (args.claude_version or args.codex_version or args.opencode_version): + if not ( + args.claude_version or args.codex_version or args.opencode_version or args.copilot_version + ): parser.error( - "Select --claude-version, --codex-version and/or --opencode-version explicitly." + "Select --claude-version, --codex-version, --opencode-version and/or " + "--copilot-version explicitly." ) if args.ug_version != "checkout" and not re.fullmatch( r"[0-9][0-9A-Za-z.!+_-]*", args.ug_version @@ -516,9 +531,11 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str: "claude": args.claude_version, "codex": args.codex_version, "opencode": args.opencode_version, + "copilot": args.copilot_version, "claude_model": args.claude_model, "codex_model": args.codex_model, "opencode_model": args.opencode_model, + "copilot_model": args.copilot_model, "claude_provider": args.claude_provider, "claude_relayed_provider": args.claude_relayed_provider, "claude_provider_model": args.claude_provider_model, @@ -731,7 +748,9 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str: report["agents"] = {} for agent in agents: agent_command = npm_executable(agent_bin, agent) - version = run([agent_command, "--version"], env=runtime_env, timeout=30) + version = agent_version_output( + agent, run([agent_command, "--version"], env=runtime_env, timeout=30) + ) expected = getattr(args, f"{agent}_version") if not re.search(rf"(? str: raise RuntimeError("Pytest returned success without a test report.") # A bootstrap/update path must not silently alter the selected agent version. for agent in agents: - after = run( - [npm_executable(agent_bin, agent), "--version"], - env=runtime_env, - timeout=30, + after = agent_version_output( + agent, + run( + [npm_executable(agent_bin, agent), "--version"], + env=runtime_env, + timeout=30, + ), ) if after != report["agents"][agent]: raise RuntimeError(f"{agent} changed version during the suite: {after}") diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 3495ab2a7..3025ac606 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -465,7 +465,7 @@ def configure_tool( result = kilo.write_tool_config(state, model) else: result = opencode.write_tool_config(state, model) - # gemini/opencode/copilot/pi return (state, token); codex/claude return state + # gemini/opencode/kilo/pi return (state, token); codex/claude/copilot return state if isinstance(result, tuple): return result[0] return result diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 40c33f386..1984d4811 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -1,10 +1,14 @@ -"""GitHub Copilot CLI agent: writes ~/.copilot/.env and injects BYOK env vars at launch. +"""GitHub Copilot CLI agent: writes ~/.copilot/ucode.env and injects BYOK env vars at launch. Copilot CLI's BYOK config (COPILOT_PROVIDER_*) is documented as env-var only — -the CLI does not auto-load ~/.copilot/.env. We still write the file so users can -inspect what's configured (`cat ~/.copilot/.env`) and to give `revert` something -to clean up; the values are also injected directly into the child process's -environment at launch. +the CLI does not auto-load ~/.copilot/ucode.env. We still write the file so users can +inspect what's configured and to give `revert` something to clean up; it never holds +a credential. + +Copilot holds its provider credential for the whole session, so `ug copilot` runs it +against a per-session loopback proxy (`gateway_proxy.start_gateway_auth_proxy`): +Copilot authenticates to the proxy with a random session secret, and the proxy +forwards each request with a fresh Databricks token. System Claude models (`system.ai.claude-*`) use Copilot CLI's `anthropic` provider against the gateway's Anthropic endpoint, with `COPILOT_PROVIDER_MODEL_ID` set to Copilot's catalog id @@ -22,9 +26,12 @@ import os import re -import signal +import secrets import threading +import time +from collections.abc import Callable from pathlib import Path +from typing import Protocol from ucode.config_io import ( APP_DIR, @@ -36,22 +43,26 @@ write_dotenv, write_json_file, ) +from ucode.constants import AI_GATEWAY_REQUEST_TAGS_HEADER, LOOPBACK_HOST +from ucode.custom_oauth import get_custom_client_token from ucode.databricks import ( - TOKEN_REFRESH_INTERVAL_SECONDS, build_copilot_base_url, build_tool_base_url, get_databricks_token, ) from ucode.os_compatibility import subprocess_cross_os -from ucode.state import mark_tool_managed, save_state +from ucode.state import is_tool_managed, mark_tool_managed, save_state +from ucode.telemetry import agent_version, ug_version from .args import LaunchOptions, explicit_model_arg_value +from .request_tags import request_tags_header_value 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_BACKUP_PATH = APP_DIR / "copilot-ucode-env.backup" COPILOT_MCP_BACKUP_PATH = APP_DIR / "copilot-ucode-mcp-config.backup.json" +COPILOT_PROXY_LOG_PATH = APP_DIR / "logs" / "copilot-proxy.log" SPEC: ToolSpec = { "binary": "copilot", @@ -67,15 +78,21 @@ "COPILOT_PROVIDER_WIRE_API", "COPILOT_PROVIDER_MODEL_ID", "COPILOT_MODEL", - "COPILOT_PROVIDER_BEARER_TOKEN", "COPILOT_OFFLINE", - "OAUTH_TOKEN", ] LEGACY_ENV_KEYS = [ "OPENAI_BASE_URL", "OPENAI_API_KEY", "COPILOT_PROVIDER_API_KEY", ] +# Earlier versions wrote the Databricks token (or PAT) under these keys. +CREDENTIAL_ENV_KEYS = ("COPILOT_PROVIDER_BEARER_TOKEN", "OAUTH_TOKEN") +# Inherited settings that would bypass the proxy's credential or its HTTP transport. +_INHERITED_PROVIDER_OVERRIDES = ( + "COPILOT_PROVIDER_API_KEY", + "COPILOT_PROVIDER_API_KEY_COMMAND", + "COPILOT_PROVIDER_TRANSPORT", +) _GPT_MODEL_MAJOR_PATTERN = re.compile(r"^(?:system\.ai\.)?(?:databricks-)?gpt-(\d+)(?=$|[.-])") @@ -130,48 +147,50 @@ def default_model(state: dict) -> str | None: def render_env_overlay( - workspace: str, + base_url: str, selected_model: str, - token: str, *, override_model: str | None = None, ) -> dict[str, str]: + """Copilot's provider settings for ``selected_model``, rooted at ``base_url``. + + ``base_url`` is the workspace for the inspectable env file and the loopback proxy + at launch. Credentials are added only to the launch environment.""" request_model = override_model or selected_model catalog_model_id = copilot_catalog_model_id(request_model) if catalog_model_id is not None: return { "COPILOT_PROVIDER_TYPE": "anthropic", - "COPILOT_PROVIDER_BASE_URL": build_tool_base_url("claude", workspace), + "COPILOT_PROVIDER_BASE_URL": build_tool_base_url("claude", base_url), "COPILOT_PROVIDER_MODEL_ID": catalog_model_id, "COPILOT_MODEL": selected_model, - "COPILOT_PROVIDER_BEARER_TOKEN": token, "COPILOT_OFFLINE": "true", - "OAUTH_TOKEN": token, } wire_api = "responses" if model_uses_responses_api(request_model) else "completions" return { "COPILOT_PROVIDER_TYPE": "openai", - "COPILOT_PROVIDER_BASE_URL": build_copilot_base_url(workspace), + "COPILOT_PROVIDER_BASE_URL": build_copilot_base_url(base_url), "COPILOT_PROVIDER_WIRE_API": wire_api, "COPILOT_MODEL": selected_model, - "COPILOT_PROVIDER_BEARER_TOKEN": token, "COPILOT_OFFLINE": "true", - "OAUTH_TOKEN": token, } -def build_runtime_env(workspace: str, model: str, token: str) -> dict[str, str]: +def build_runtime_env(proxy_base_url: str, model: str, session_secret: str) -> dict[str, str]: + """Copilot's launch environment: provider settings pointed at the session proxy.""" env = os.environ.copy() + for key in _INHERITED_PROVIDER_OVERRIDES: + env.pop(key, None) override_model = env.get("COPILOT_PROVIDER_WIRE_MODEL") - env.update(render_env_overlay(workspace, model, token, override_model=override_model)) + env.update(render_env_overlay(proxy_base_url, model, override_model=override_model)) + env["COPILOT_PROVIDER_BEARER_TOKEN"] = session_secret return env def build_mcp_server_entry(argv: list[str]) -> dict: # A `local` MCP server runs a stdio command; `command`/`args` split the # argv. ug registers the `ug mcp-proxy ...` bridge here so Copilot - # never speaks HTTP+bearer directly — the proxy handles token refresh. The - # OAUTH_TOKEN env Copilot still injects at launch is for MODEL auth, not MCP. + # never speaks HTTP+bearer directly — the proxy handles token refresh. return { "type": "local", "command": argv[0], @@ -211,22 +230,30 @@ def write_user_mcp_servers(add: dict[str, dict], remove: set[str]) -> set[str]: ) -def write_tool_config( - state: dict, - model: str, - token: str | None = None, - *, - force_refresh: bool = False, -) -> tuple[dict, str]: - backup_existing_file(COPILOT_ENV_PATH, COPILOT_BACKUP_PATH) - if token is None: - token = get_databricks_token( - state["workspace"], state.get("profile"), force_refresh=force_refresh - ) +def scrub_stored_credentials() -> None: + """Remove tokens earlier versions persisted, so neither the env file nor `ug revert` keeps one. + + A backup holding a credential key was produced by ug itself (users don't own + ucode.env), so it is dropped rather than restored.""" + if COPILOT_BACKUP_PATH.exists() and any( + key in parse_dotenv(COPILOT_BACKUP_PATH) for key in CREDENTIAL_ENV_KEYS + ): + COPILOT_BACKUP_PATH.unlink() + existing = parse_dotenv(COPILOT_ENV_PATH) + if any(key in existing for key in CREDENTIAL_ENV_KEYS): + for key in CREDENTIAL_ENV_KEYS: + existing.pop(key, None) + write_dotenv(COPILOT_ENV_PATH, existing) + + +def write_tool_config(state: dict, model: str) -> dict: + scrub_stored_credentials() + if not is_tool_managed(state, "copilot"): + backup_existing_file(COPILOT_ENV_PATH, COPILOT_BACKUP_PATH) existing = parse_dotenv(COPILOT_ENV_PATH) # Keep the inspectable file self-consistent without treating it as launch input. override_model = existing.get("COPILOT_PROVIDER_WIRE_MODEL") - overlay = render_env_overlay(state["workspace"], model, token, override_model=override_model) + overlay = render_env_overlay(state["workspace"], model, override_model=override_model) for key in LEGACY_ENV_KEYS: existing.pop(key, None) # Wire API is openai-only and the catalog id is anthropic-only; drop whichever this model doesn't use. @@ -237,53 +264,100 @@ def write_tool_config( write_dotenv(COPILOT_ENV_PATH, existing) state = mark_tool_managed(state, "copilot", MANAGED_KEYS) save_state(state) - return state, token + return state + + +def _token_provider(state: dict) -> Callable[[bool], str]: + """Mint from the source the rest of ug uses for this workspace: custom OAuth when configured, + else the CLI profile.""" + workspace = state["workspace"] + custom_oauth = state.get("custom_oauth") + + def provider(force_refresh: bool) -> str: + if isinstance(custom_oauth, dict): + return get_custom_client_token( + workspace, + custom_oauth["client_id"], + custom_oauth["redirect_url"], + scopes=custom_oauth["scopes"], + profile=custom_oauth.get("profile"), + force_refresh=force_refresh, + ) + return get_databricks_token(workspace, state.get("profile"), force_refresh=force_refresh) + + return provider + + +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": f"ucode/{ug_version()} copilot/{agent_version(SPEC['binary'])}"} + request_tags = request_tags_header_value() + if request_tags: + headers[AI_GATEWAY_REQUEST_TAGS_HEADER] = request_tags + managed = state.get("copilot_http_headers") + for name, value in (managed if isinstance(managed, dict) else {}).items(): + if not isinstance(name, str) or not isinstance(value, str): + continue + if name.strip().casefold() == "authorization": + continue + for existing in [key for key in headers if key.casefold() == name.strip().casefold()]: + headers.pop(existing) + headers[name.strip()] = value + return headers -def _refresh_token_once( - state: dict, - model: str | None = None, - *, - force_refresh: bool = False, -) -> tuple[str, str]: - model = model or default_model(state) - if not model: - raise RuntimeError("No Copilot model is available on this workspace.") - _, token = write_tool_config(state, model, force_refresh=force_refresh) - return model, token +def _log_proxy_failure(exc: BaseException) -> None: + # Copilot owns the terminal, so failures go to a file instead of stderr. + try: + COPILOT_PROXY_LOG_PATH.parent.mkdir(parents=True, exist_ok=True) + with COPILOT_PROXY_LOG_PATH.open("a", encoding="utf-8") as log: + log.write(f"{time.strftime('%Y-%m-%dT%H:%M:%S%z')} token refresh failed: {exc}\n") + except OSError: + pass + +class _ChildProcess(Protocol): + def wait(self) -> int: ... -def _refresh_forever(state: dict, model: str, stop_event: threading.Event) -> None: - while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS): + +def _wait_for_exit(proc: _ChildProcess) -> int: + while True: try: - _refresh_token_once(state, model, force_refresh=True) - except RuntimeError: + return proc.wait() + except KeyboardInterrupt: + # The terminal delivers Ctrl-C to Copilot too; keep the proxy up until it exits. continue def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: model = explicit_model_arg_value(tool_args) or options.user_pinned_model or default_model(state) - model, token = _refresh_token_once(state, model) - env = build_runtime_env(state["workspace"], model, token) - - stop_event = threading.Event() - refresher = threading.Thread( - target=_refresh_forever, - args=(state, model, stop_event), - daemon=True, + if not model: + raise RuntimeError("No Copilot model is available on this workspace.") + scrub_stored_credentials() + + # Only Copilot launches need the proxy (and its `httpx` import). + from ucode import gateway_proxy + + workspace = state["workspace"] + session_secret = secrets.token_urlsafe(32) + server, cache, client = gateway_proxy.start_gateway_auth_proxy( + workspace, + _token_provider(state), + session_secret=session_secret, + added_headers=proxy_headers(state), + sign_in_hint=f"Run `databricks auth login --host {workspace}`, then restart `ug copilot`.", + on_refresh_failure=_log_proxy_failure, ) - refresher.start() - - proc = subprocess_cross_os.popen([SPEC["binary"], *mcp_config_args(), *tool_args], env=env) + threading.Thread(target=server.serve_forever, daemon=True).start() try: - returncode = proc.wait() - except KeyboardInterrupt: - proc.send_signal(signal.SIGINT) - returncode = proc.wait() + proxy_base_url = f"http://{LOOPBACK_HOST}:{server.server_address[1]}" + env = build_runtime_env(proxy_base_url, model, session_secret) + proc = subprocess_cross_os.popen([SPEC["binary"], *mcp_config_args(), *tool_args], env=env) + returncode = _wait_for_exit(proc) finally: - stop_event.set() - refresher.join(timeout=1) - + cache.stop() + server.shutdown() + client.close() raise SystemExit(returncode) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 055143a89..7367890a9 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -39,6 +39,7 @@ ) from ucode.agents import claude as claude_agent from ucode.agents import codex as codex_agent +from ucode.agents import copilot as copilot_agent from ucode.agents import ( launch as launch_agent, ) @@ -1393,6 +1394,8 @@ def revert() -> int: desktop_result = revert_desktop() managed_configs = state.get("managed_configs") or {} mcp_results = revert_mcp_configs(state) + # Before restoring backups: one written by an earlier ug may hold a token. + copilot_agent.scrub_stored_credentials() claude_managed_result = claude_agent.revert_managed_settings() codex_managed_result = codex_agent.revert_managed_config() @@ -3509,7 +3512,7 @@ def _launch_tool( f"{TOOL_SPECS[tool]['display']} may require one-time hook review. Open " "`/hooks` and trust the ug routing hooks if prompted." ) - if tool in ("gemini", "opencode", "kilo", "copilot", "pi"): + if tool in ("gemini", "opencode", "kilo", "pi"): print_note( f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically " f"every 30 minutes while the session is running." diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 7c2e83584..74fc1e1f4 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -12,6 +12,10 @@ subscription models keep the OAuth passthrough, so one Claude Code session can use both. +The gateway-auth proxy serves agents that can only hold a static credential +(Copilot CLI): the agent authenticates to the proxy with a per-session secret, +and the proxy swaps it for a fresh Databricks token on allowlisted gateway paths. + Security invariants (mirroring `databricks.py` token handling): - Binds 127.0.0.1 only; never exposed off-host. - Never logs header values or bodies. The Databricks token lives in memory, @@ -23,15 +27,17 @@ import base64 import binascii +import hmac import json import os import sys import threading import time import uuid -from collections.abc import Callable +from collections.abc import Callable, Mapping from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import cast +from urllib.parse import unquote import httpx @@ -127,12 +133,15 @@ def __init__( token_provider: Callable[[bool], str], *, force_refresh_near_expiry: bool = False, + on_refresh_failure: Callable[[BaseException], None] = log_token_refresh_failure, ) -> None: # token_provider(force_refresh) mints a fresh token from the same source the # client agent uses, so the proxy authenticates as the same principal (e.g. # a per-user custom-OAuth token, not the default CLI profile). self._token_provider = token_provider self._force_refresh_near_expiry = force_refresh_near_expiry + # Agents that own the terminal (a TUI) need failures kept off stderr. + self._on_refresh_failure = on_refresh_failure self._state_lock = threading.Lock() # guards _token / _expiry (brief) self._refresh_lock = threading.Lock() # single-flights the CLI refresh self._stop = threading.Event() @@ -165,7 +174,7 @@ def _ensure_fresh(self) -> None: except RuntimeError as exc: # Keep serving the current token; a request that then 401s triggers # a forced refresh + retry (see _ProxyHandler._handle). - log_token_refresh_failure(exc) + self._on_refresh_failure(exc) @property def token(self) -> str: @@ -173,6 +182,9 @@ def token(self) -> str: with self._state_lock: return self._token + def report_refresh_failure(self, exc: BaseException) -> None: + self._on_refresh_failure(exc) + def refresh(self) -> None: """Force a fresh mint now (used by the retry-on-401 path).""" with self._refresh_lock: @@ -185,7 +197,7 @@ def run_refresher(self) -> None: except Exception as exc: # noqa: BLE001 - a stray error must NOT kill the thread # If this thread dies, nothing refreshes and the session lapses at # the ~1h mark until restart. Log and keep looping instead. - log_token_refresh_failure(exc) + self._on_refresh_failure(exc) def stop(self) -> None: self._stop.set() @@ -257,6 +269,14 @@ def _safe_send_error(self, code: int, message: str) -> None: def _forward_target(self, body: bytes | None) -> tuple[str, frozenset[str], str]: return self.token_header, frozenset(), "forward" + def _added_headers(self) -> dict[str, str]: + return {} + + def _on_refresh_failure(self, exc: RuntimeError) -> bool: + """Report a failed forced refresh; True means a response was already sent.""" + log_token_refresh_failure(exc) + return False + def _handle(self) -> None: diagnostic_id = uuid.uuid4().hex[:12] started = time.monotonic() @@ -273,9 +293,11 @@ def _handle(self) -> None: ) def request_headers() -> dict[str, str]: - return forwarded_request_headers( + headers = forwarded_request_headers( self, self.cache.token, token_header, extra_strip=extra_strip ) + headers.update(self._added_headers()) + return headers try: # First attempt with the current token. @@ -308,7 +330,8 @@ def request_headers() -> dict[str, str]: # `databricks auth login` hint rather than silently relaying a bare 401, # which otherwise reads as an Anthropic `/login` prompt and sends the # user to the wrong re-auth. Still retry + relay with the existing token. - log_token_refresh_failure(exc) + if self._on_refresh_failure(exc): + return headers = request_headers() with self.client.stream(self.command, url, headers=headers, content=body) as resp: log_proxy_diagnostic( @@ -428,6 +451,98 @@ def _forward_target(self, body: bytes | None) -> tuple[str, frozenset[str], str] return self.token_header, frozenset(), "relay" +# Gateway model routes a gateway-auth proxy session may reach. Everything else +# (workspace REST APIs included) is refused, so the session secret only buys inference. +GATEWAY_MODEL_PATH_PREFIXES = ( + "/ai-gateway/anthropic/", + "/ai-gateway/mlflow/v1/", + "/ai-gateway/openai/v1/", +) +_ANTHROPIC_PATH_PREFIX = "/ai-gateway/anthropic/" +# The agent's credential headers carry the session secret, never a Databricks token. +_SESSION_CREDENTIAL_HEADERS = frozenset({"authorization", "x-api-key"}) + + +def _is_allowed_gateway_path(path: str, prefixes: tuple[str, ...]) -> bool: + decoded = path.split("?", 1)[0] + # Decode until stable so a double-encoded dot segment can't reach another API. + while (next_decoded := unquote(decoded)) != decoded: + decoded = next_decoded + if any(segment in (".", "..") for segment in decoded.split("/")): + return False + return decoded.startswith(prefixes) + + +def _bearer_value(value: str | None) -> str | None: + if not value: + return None + scheme, _, credential = value.partition(" ") + return credential.strip() if scheme.lower() == "bearer" else None + + +class _GatewayAuthProxyHandler(_ProxyHandler): + # Set by the server factory. + session_secret: str + added_headers: Mapping[str, str] + allowed_path_prefixes: tuple[str, ...] + sign_in_hint: str + + def _has_session_secret(self) -> bool: + expected = self.session_secret.encode() + for presented in ( + _bearer_value(self.headers.get(AUTHORIZATION_HEADER)), + self.headers.get("x-api-key"), + ): + if presented and hmac.compare_digest(presented.encode(), expected): + return True + return False + + def _reject(self, status: int, message: str) -> None: + length = int(self.headers.get("Content-Length", 0) or 0) + if length: + self.rfile.read(length) + self._send_json_error(status, message) + + def _send_json_error(self, status: int, message: str) -> None: + # Shaped like the upstream dialect so the agent shows the message instead of + # a generic transport failure. + if self.path.startswith(_ANTHROPIC_PATH_PREFIX): + error_type = "authentication_error" if status == 401 else "not_found_error" + payload = {"type": "error", "error": {"type": error_type, "message": message}} + else: + payload = {"error": {"message": message, "type": "invalid_request_error"}} + body = json.dumps(payload).encode() + try: + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except OSError: + pass + + def _handle(self) -> None: + if not _is_allowed_gateway_path(self.path, self.allowed_path_prefixes): + self._reject(404, "This path is not served by the Unity Gateway proxy.") + return + if not self._has_session_secret(): + self._reject(401, "Missing or invalid Unity Gateway session credential.") + return + super()._handle() + + def _forward_target(self, body: bytes | None) -> tuple[str, frozenset[str], str]: + added = frozenset(name.lower() for name in self.added_headers) + return self.token_header, _SESSION_CREDENTIAL_HEADERS | added, "gateway" + + def _added_headers(self) -> dict[str, str]: + return dict(self.added_headers) + + def _on_refresh_failure(self, exc: RuntimeError) -> bool: + self.cache.report_refresh_failure(exc) + self._send_json_error(401, f"Databricks sign-in expired. {self.sign_in_hint}") + return True + + class _LoopbackHTTPServer(ThreadingHTTPServer): # Windows permits sharing an active listener when SO_REUSEADDR is enabled. allow_reuse_address = ThreadingHTTPServer.allow_reuse_address and os.name != "nt" @@ -443,6 +558,8 @@ def _start_proxy( handler_type: type[_ProxyHandler], force_refresh_near_expiry: bool, fallback_on_bind_error: bool = True, + handler_attrs: Mapping[str, object] | None = None, + on_refresh_failure: Callable[[BaseException], None] = log_token_refresh_failure, ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -457,7 +574,11 @@ def _start_proxy( thread) and calls shutdown()/cache.stop()/client.close() on exit. """ upstream_base = f"{workspace.rstrip('/')}/{upstream_path.lstrip('/')}" - cache = TokenCache(token_provider, force_refresh_near_expiry=force_refresh_near_expiry) + cache = TokenCache( + token_provider, + force_refresh_near_expiry=force_refresh_near_expiry, + on_refresh_failure=on_refresh_failure, + ) # One pooled, keep-alive client shared across handler threads: reuses TCP+TLS # to the gateway instead of a fresh handshake per request. Don't follow # redirects — a proxy relays 3xx verbatim. @@ -472,6 +593,7 @@ def _start_proxy( "cache": cache, "client": client, "token_header": token_header, + **(handler_attrs or {}), }, ), ) @@ -525,3 +647,37 @@ def start_otel_proxy( force_refresh_near_expiry=True, fallback_on_bind_error=False, ) + + +def start_gateway_auth_proxy( + workspace: str, + token_provider: Callable[[bool], str], + *, + session_secret: str, + added_headers: Mapping[str, str], + sign_in_hint: str, + on_refresh_failure: Callable[[BaseException], None] = log_token_refresh_failure, +) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: + """Start a per-session proxy for agents that only hold a static credential. + + The agent sends ``session_secret`` as its API key; the proxy accepts it only on + gateway model paths and forwards with a fresh Databricks token plus + ``added_headers``. Binds an OS-assigned port, read from ``server.server_address[1]``. + """ + return _start_proxy( + workspace, + token_provider, + 0, + upstream_path="", + token_header=AUTHORIZATION_HEADER, + handler_type=_GatewayAuthProxyHandler, + force_refresh_near_expiry=True, + fallback_on_bind_error=False, + handler_attrs={ + "session_secret": session_secret, + "added_headers": dict(added_headers), + "allowed_path_prefixes": GATEWAY_MODEL_PATH_PREFIXES, + "sign_in_hint": sign_in_hint, + }, + on_refresh_failure=on_refresh_failure, + ) diff --git a/tests/AGENTS.md b/tests/AGENTS.md index ad1618979..8ff486fbd 100644 --- a/tests/AGENTS.md +++ b/tests/AGENTS.md @@ -92,7 +92,8 @@ exercise. and Docker build files belong in `integration/utils/`; keep pytest entry points and run documentation at the suite root. - Current scope is basic Claude/Codex configuration, routing, script usage, - command forwarding, and configure/revert CUJs. Do not add MCP/skills functionality, + command forwarding, and configure/revert CUJs, plus Copilot CLI headless + model tasks (non-blocking lane). Do not add 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 d6eb889cc..b01dd5cea 100644 --- a/tests/README.md +++ b/tests/README.md @@ -432,6 +432,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_case_14_*` | Launch configured and fresh Codex with a model location | The app-server list exactly matches the independent API-compatible parent catalog, includes the dedicated Codex service, and contains no out-of-schema models | | `test_ug_claude_headless_prompt_argument`, `test_ug_claude_headless_prompt_stdin`, `test_ug_claude_headless_prompt_after_separator` | Run Claude from a script using each prompt form | Structured final answer contains the file value; exit zero; no routing | | `test_ug_codex_headless_prompt_argument`, `test_ug_codex_headless_prompt_stdin`, `test_ug_codex_headless_prompt_after_separator` | Run Codex from a script using each prompt form | Completed turn and final answer contain the file value; exit zero; no routing | +| `test_ug_copilot_headless_completes_file_task`, `test_ug_copilot_headless_gpt_responses_task` | Run Copilot CLI from a script (`-p ... --output-format json`) on a Claude model (Anthropic Messages) and a GPT model (Responses) | Completed turn whose assistant message contains the file value; no file under the session home holds the Databricks credential (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 | | `test_ug_agents_admin_managed_guardrails` | `ug agents add` / `remove` on an agent the admin config enables | Add is a no-op noting the admin manages it; remove is rejected with nonzero exit | diff --git a/tests/integration/README.md b/tests/integration/README.md index 1345f5b7d..fada250fd 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -203,6 +203,12 @@ when an explicit model is supplied. The only live OpenCode journey is the headle prompt case in `test_ug_opencode_headless.py`; there is no dedicated live OpenCode model-selection journey. +Copilot CLI has two live headless journeys in `test_ug_copilot_headless.py`, one per +wire API (Anthropic Messages for Claude, Responses for GPT). Both go through the +session's loopback proxy and check that no file under the session home holds the +Databricks credential. Token refresh past expiry is covered by the hermetic proxy +tests in `tests/test_gateway_proxy_integration.py`, not by a live hour-long session. + ## Shared agent helpers `utils/agents/claude.py` and `utils/agents/codex.py` own the existing native answer @@ -256,6 +262,7 @@ exercise explicit model arguments use a real `system.ai` model already discovere by `ug configure`, recorded in that case's `model.json`. Optional `--claude-model` `--codex-model` and `--opencode-model` overrides reproduce a particular model-related failure. OpenCode is opt-in: pass `--opencode-version` (for example `1.18.31`) to install and select it. +Copilot CLI is opt-in the same way: pass `--copilot-version` (for example `1.0.91`). ```bash # Constrain the suspected dependency while keeping the real CLI and gateway. @@ -414,6 +421,7 @@ test_ug_claude_headless.py # script prompts, models, caller setting test_ug_claude_relayed.py # relayed session: subscription + Databricks-hosted models test_ug_claude_tracing.py # Claude OTLP export reaches the configured trace table test_ug_codex_headless.py # script prompts and model arguments +test_ug_copilot_headless.py # Copilot via the session proxy: Claude + GPT tasks, no token on disk test_ug_codex_tracing.py # Codex OTLP export reaches the configured trace table test_ug_claude_commands.py # command help forwarding test_ug_codex_commands.py # command help and parser error forwarding @@ -607,7 +615,7 @@ cases instead complete real file tasks through `--workspace` and `--model-locati journeys check setup and launch with existing services, without inference. There are **85 live cases** (including 11 marked TUI journeys) and **7 installation -checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases. One +checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds two live headless cases. One **`workspace_switch` case** uses two real workspaces and checks skills MCP cleanup and a completed Claude task. A further **43 `managed_fixture` cases** (ten of them also `live`) run on the managed workspace with a checked-in JSON CodingAgentConfig from `tests/fixtures/managed_config/` @@ -833,6 +841,8 @@ 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`, two cases) runs on Linux and Windows the +same way and is likewise outside `cujs` until it is stable. Each lane installs only the agent CLIs it needs (one, except the managed OpenCode lane's three), once, and runs all its configure, headless, diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index d149fb914..e44342629 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -27,7 +27,8 @@ def pytest_collection_modifyitems(config, items): if "monkeypatch" in item.fixturenames: raise pytest.UsageError("Use the integration runner; unit fixtures were inherited.") if any( - item.get_closest_marker(a) and a not in agents for a in ("claude", "codex", "opencode") + item.get_closest_marker(a) and a not in agents + for a in ("claude", "codex", "opencode", "copilot") ): deselected.append(item) else: diff --git a/tests/integration/pytest.ini b/tests/integration/pytest.ini index 0ee7a1694..f08989f71 100644 --- a/tests/integration/pytest.ini +++ b/tests/integration/pytest.ini @@ -12,3 +12,4 @@ markers = claude: only runs when Claude Code is explicitly selected codex: only runs when Codex is explicitly selected opencode: only runs when OpenCode is explicitly selected + copilot: only runs when GitHub Copilot CLI is explicitly selected diff --git a/tests/integration/test_ug_copilot_headless.py b/tests/integration/test_ug_copilot_headless.py new file mode 100644 index 000000000..1527d8f19 --- /dev/null +++ b/tests/integration/test_ug_copilot_headless.py @@ -0,0 +1,99 @@ +"""CUJs for using GitHub Copilot CLI from scripts through installed ug.""" + +import pytest +from utils.constants import CLAUDE_TEST_MODEL, CODEX_TEST_MODEL +from utils.evidence import FileTask + +pytestmark = [pytest.mark.live, pytest.mark.copilot] + +COPILOT_HEADLESS_FLAGS = ( + "--allow-all-tools", + "--output-format", + "json", + "--no-auto-update", + "--disable-builtin-mcps", +) + + +def _assert_no_credential_on_disk(session): + """No Databricks credential in any file ug or Copilot wrote under the session home.""" + bearer = session.env.get("DATABRICKS_BEARER") + assert bearer, "The live suite authenticates with DATABRICKS_BEARER." + for root in (session.home / ".copilot", session.home / ".ucode"): + for path in root.rglob("*") if root.exists() else (): + if path.is_file(): + text = path.read_text(encoding="utf-8", errors="replace") + assert bearer not in text, f"{path} holds the Databricks credential" + env_file = (session.home / ".copilot/ucode.env").read_text(encoding="utf-8") + assert "COPILOT_PROVIDER_BEARER_TOKEN" not in env_file + assert "OAUTH_TOKEN" not in env_file + + +def test_ug_copilot_headless_completes_file_task(live_session, workspace): + """Scenario: configure Copilot, then run a headless Claude prompt through `ug copilot`. + + Expected: Copilot reads the fixture through the session's loopback proxy (Anthropic + provider) and its completed answer contains the unknown value; no file under the + session home holds the Databricks credential afterwards. + """ + session = live_session + task = FileTask(session) + session.run( + "configure", + "--agents", + "copilot", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + session.record("copilot-model.json", {"model": CLAUDE_TEST_MODEL}) + + result = session.run( + "copilot", + "--model", + CLAUDE_TEST_MODEL, + "-p", + task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + + task.assert_headless_answer("copilot", result) + _assert_no_credential_on_disk(session) + + +def test_ug_copilot_headless_gpt_responses_task(live_session, workspace): + """Scenario: configure Copilot, then run a headless GPT prompt through `ug copilot`. + + Expected: Copilot reaches the gateway's Responses API through the session proxy and + its completed answer contains the fixture's unknown value; no file under the session + home holds the Databricks credential afterwards. + """ + session = live_session + task = FileTask(session) + session.run( + "configure", + "--agents", + "copilot", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + session.record("copilot-model.json", {"model": CODEX_TEST_MODEL}) + + result = session.run( + "copilot", + "--model", + CODEX_TEST_MODEL, + "-p", + task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + + task.assert_headless_answer("copilot", result) + _assert_no_credential_on_disk(session) diff --git a/tests/integration/utils/evidence.py b/tests/integration/utils/evidence.py index a7899bfe2..f0fba4e38 100644 --- a/tests/integration/utils/evidence.py +++ b/tests/integration/utils/evidence.py @@ -215,6 +215,15 @@ def assert_headless_answer(self, agent: str, result) -> None: if row.get("type") == "text" ) assert self.value in answer, result.stdout + elif agent == "copilot": + assert result.returncode == 0, result.stdout + assert any(row.get("type") == "assistant.turn_end" for row in payloads), result.stdout + answers = [ + (row.get("data") or {}).get("content", "") + for row in payloads + if row.get("type") == "assistant.message" + ] + assert any(self.value in answer for answer in answers), result.stdout else: raise AssertionError(f"Unknown agent: {agent}") diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 430b3563e..97359c0dd 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -3,8 +3,8 @@ from __future__ import annotations import json -import threading +import httpx import pytest from ucode.agents import copilot @@ -28,11 +28,11 @@ def test_config_path_is_ucode_env_file(self): class TestRenderEnvOverlay: def test_sets_provider_base_url(self): - env = copilot.render_env_overlay(WS, "gpt-5", "tok") + env = copilot.render_env_overlay(WS, "gpt-5") assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" def test_claude_uses_anthropic_provider_with_catalog_model_id(self): - env = copilot.render_env_overlay(WS, "system.ai.claude-sonnet-5-5", "tok") + env = copilot.render_env_overlay(WS, "system.ai.claude-sonnet-5-5") assert env["COPILOT_PROVIDER_TYPE"] == "anthropic" assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/anthropic" @@ -41,27 +41,27 @@ def test_claude_uses_anthropic_provider_with_catalog_model_id(self): assert "COPILOT_PROVIDER_WIRE_API" not in env def test_claude_wire_model_override_selects_anthropic_provider(self): - env = copilot.render_env_overlay( - WS, "gpt-5", "tok", override_model="system.ai.claude-opus-4-7" - ) + env = copilot.render_env_overlay(WS, "gpt-5", override_model="system.ai.claude-opus-4-7") assert env["COPILOT_PROVIDER_TYPE"] == "anthropic" assert env["COPILOT_PROVIDER_MODEL_ID"] == "claude-opus-4.7" def test_sets_provider_type(self): - env = copilot.render_env_overlay(WS, "m", "t") + env = copilot.render_env_overlay(WS, "m") assert env["COPILOT_PROVIDER_TYPE"] == "openai" def test_sets_model(self): - env = copilot.render_env_overlay(WS, "claude-sonnet-4-6", "tok") + env = copilot.render_env_overlay(WS, "claude-sonnet-4-6") assert env["COPILOT_MODEL"] == "claude-sonnet-4-6" - def test_sets_bearer_token(self): - env = copilot.render_env_overlay(WS, "m", "tok123") - assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "tok123" + @pytest.mark.parametrize("model", ["m", "system.ai.claude-sonnet-5-5"]) + def test_carries_no_credential(self, model): + env = copilot.render_env_overlay(WS, model) + for key in copilot.CREDENTIAL_ENV_KEYS: + assert key not in env def test_sets_offline_true(self): - env = copilot.render_env_overlay(WS, "m", "t") + env = copilot.render_env_overlay(WS, "m") assert env["COPILOT_OFFLINE"] == "true" @pytest.mark.parametrize( @@ -78,25 +78,46 @@ def test_sets_offline_true(self): def test_selects_wire_api_from_override_and_keeps_selected_model( self, selected_model, override_model, expected_api ): - env = copilot.render_env_overlay(WS, selected_model, "tok", override_model=override_model) + env = copilot.render_env_overlay(WS, selected_model, override_model=override_model) assert env["COPILOT_PROVIDER_WIRE_API"] == expected_api assert env["COPILOT_MODEL"] == selected_model +PROXY = "http://127.0.0.1:4321" + + class TestBuildRuntimeEnv: def test_inherits_path(self): - env = copilot.build_runtime_env(WS, "m", "t") + env = copilot.build_runtime_env(PROXY, "m", "secret") assert "PATH" in env - def test_overrides_copilot_vars(self): - env = copilot.build_runtime_env(WS, "m", "tok") - assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" - assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "tok" + def test_points_copilot_at_the_session_proxy(self): + env = copilot.build_runtime_env(PROXY, "m", "secret") + assert env["COPILOT_PROVIDER_BASE_URL"] == f"{PROXY}/ai-gateway/mlflow/v1" + assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "secret" - def test_sets_oauth_token_for_mcp(self): - env = copilot.build_runtime_env(WS, "m", "tok") - assert env["OAUTH_TOKEN"] == "tok" + def test_claude_goes_through_the_proxy_anthropic_route(self): + env = copilot.build_runtime_env(PROXY, "system.ai.claude-sonnet-5-5", "secret") + assert env["COPILOT_PROVIDER_BASE_URL"] == f"{PROXY}/ai-gateway/anthropic" + + def test_sets_no_oauth_token(self, monkeypatch): + monkeypatch.delenv("OAUTH_TOKEN", raising=False) + env = copilot.build_runtime_env(PROXY, "m", "secret") + assert "OAUTH_TOKEN" not in env + + @pytest.mark.parametrize( + "key", + [ + "COPILOT_PROVIDER_API_KEY", + "COPILOT_PROVIDER_API_KEY_COMMAND", + "COPILOT_PROVIDER_TRANSPORT", + ], + ) + def test_drops_inherited_settings_that_bypass_the_proxy(self, monkeypatch, key): + monkeypatch.setenv(key, "user-value") + env = copilot.build_runtime_env(PROXY, "m", "secret") + assert key not in env def test_inherited_wire_model_overrides_route_without_being_cleared(self, monkeypatch): monkeypatch.setenv("COPILOT_PROVIDER_WIRE_MODEL", "gpt-4.1") @@ -104,7 +125,7 @@ def test_inherited_wire_model_overrides_route_without_being_cleared(self, monkey monkeypatch.setenv("COPILOT_PROVIDER_MODEL_ID", "user-model-id") monkeypatch.setenv("COPILOT_PROVIDER_MODEL_LIMITS_ID", "user-model-limits") - env = copilot.build_runtime_env(WS, "system.ai.gpt-6-astra", "tok") + env = copilot.build_runtime_env(PROXY, "system.ai.gpt-6-astra", "secret") assert env["COPILOT_PROVIDER_WIRE_MODEL"] == "gpt-4.1" assert env["COPILOT_PROVIDER_WIRE_API"] == "completions" @@ -326,12 +347,9 @@ def test_manages_wire_api_and_preserves_wire_model_and_user_values(self, tmp_pat ) monkeypatch.setenv("COPILOT_PROVIDER_WIRE_MODEL", "gpt-5") - state, token = copilot.write_tool_config( - {"workspace": WS}, "system.ai.gpt-5-6-sol", token="tok" - ) + state = copilot.write_tool_config({"workspace": WS}, "system.ai.gpt-5-6-sol") written = copilot.parse_dotenv(env_path) - assert token == "tok" assert written["COPILOT_MODEL"] == "system.ai.gpt-5-6-sol" assert written["COPILOT_PROVIDER_WIRE_API"] == "responses" assert written["COPILOT_PROVIDER_WIRE_MODEL"] == "system.ai.gpt-6-1-sol" @@ -346,13 +364,184 @@ def test_claude_drops_stale_wire_api_from_env_file(self, tmp_path, monkeypatch): env_path.parent.mkdir(parents=True) env_path.write_text("COPILOT_PROVIDER_WIRE_API=completions\n", encoding="utf-8") - copilot.write_tool_config({"workspace": WS}, "system.ai.claude-sonnet-5-5", token="tok") + copilot.write_tool_config({"workspace": WS}, "system.ai.claude-sonnet-5-5") written = copilot.parse_dotenv(env_path) assert "COPILOT_PROVIDER_WIRE_API" not in written assert written["COPILOT_PROVIDER_TYPE"] == "anthropic" assert written["COPILOT_PROVIDER_MODEL_ID"] == "claude-sonnet-5.5" + def test_needs_no_token_and_writes_none(self, tmp_path, monkeypatch): + env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) + + def no_token(*args, **kwargs): + raise AssertionError("configure must not mint a token") + + monkeypatch.setattr(copilot, "get_databricks_token", no_token) + + copilot.write_tool_config({"workspace": WS}, "gpt-5") + + written = copilot.parse_dotenv(env_path) + assert written["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" + for key in copilot.CREDENTIAL_ENV_KEYS: + assert key not in written + + def test_scrubs_a_token_an_earlier_version_wrote(self, tmp_path, monkeypatch): + env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) + env_path.parent.mkdir(parents=True) + env_path.write_text( + 'COPILOT_PROVIDER_BEARER_TOKEN="old-token"\nOAUTH_TOKEN="old-token"\n', + encoding="utf-8", + ) + + copilot.write_tool_config({"workspace": WS}, "gpt-5") + + assert "old-token" not in env_path.read_text(encoding="utf-8") + + def test_does_not_back_up_its_own_file_on_reconfigure(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + + state = copilot.write_tool_config({"workspace": WS}, "gpt-5") + copilot.write_tool_config(state, "gpt-5") + + assert not copilot.COPILOT_BACKUP_PATH.exists() + + +class TestScrubStoredCredentials: + def test_drops_a_token_bearing_backup_so_revert_cannot_restore_it(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + backup = copilot.COPILOT_BACKUP_PATH + backup.parent.mkdir(parents=True, exist_ok=True) + backup.write_text('COPILOT_PROVIDER_BEARER_TOKEN="old-token"\n', encoding="utf-8") + + copilot.scrub_stored_credentials() + + assert not backup.exists() + + def test_keeps_a_backup_without_credentials(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + backup = copilot.COPILOT_BACKUP_PATH + backup.parent.mkdir(parents=True, exist_ok=True) + backup.write_text('USER_SETTING="keep"\n', encoding="utf-8") + + copilot.scrub_stored_credentials() + + assert backup.read_text(encoding="utf-8") == 'USER_SETTING="keep"\n' + + def test_leaves_a_credential_free_env_file_untouched(self, tmp_path, monkeypatch): + env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) + env_path.parent.mkdir(parents=True) + env_path.write_text("COPILOT_MODEL=gpt-5\n", encoding="utf-8") + + copilot.scrub_stored_credentials() + + assert env_path.read_text(encoding="utf-8") == "COPILOT_MODEL=gpt-5\n" + + +class TestProxyHeaders: + def test_identifies_ucode_and_copilot(self, monkeypatch): + monkeypatch.setattr(copilot, "agent_version", lambda binary: "1.0.91") + monkeypatch.setattr(copilot, "ug_version", lambda: "0.9.0") + + headers = copilot.proxy_headers({}) + + assert headers["User-Agent"] == "ucode/0.9.0 copilot/1.0.91" + + def test_adds_request_tags(self, monkeypatch): + monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") + monkeypatch.setenv("AI_GATEWAY_REQUEST_TAGS", '{"team":"ug"}') + + headers = copilot.proxy_headers({}) + + assert headers["Databricks-Ai-Gateway-Request-Tags"] == '{"team":"ug"}' + + def test_managed_headers_win_by_name_but_never_replace_authorization(self, monkeypatch): + monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") + monkeypatch.delenv("AI_GATEWAY_REQUEST_TAGS", raising=False) + + headers = copilot.proxy_headers( + { + "copilot_http_headers": { + "user-agent": "admin-agent", + "X-Team": "data", + "Authorization": "Bearer admin", + } + } + ) + + assert headers == {"user-agent": "admin-agent", "X-Team": "data"} + + +class TestTokenProvider: + def test_uses_the_cli_profile(self, monkeypatch): + calls = [] + + def get_token(workspace, profile=None, *, force_refresh=False): + calls.append((workspace, profile, force_refresh)) + return "profile-token" + + monkeypatch.setattr(copilot, "get_databricks_token", get_token) + + token = copilot._token_provider({"workspace": WS, "profile": "dev"})(True) + + assert token == "profile-token" + assert calls == [(WS, "dev", True)] + + def test_uses_custom_oauth_when_configured(self, monkeypatch): + calls = [] + + def get_custom(workspace, client_id, redirect_url, *, scopes, profile, force_refresh): + calls.append((workspace, client_id, redirect_url, scopes, profile, force_refresh)) + return "custom-token" + + monkeypatch.setattr(copilot, "get_custom_client_token", get_custom) + state = { + "workspace": WS, + "custom_oauth": { + "client_id": "cid", + "redirect_url": "http://localhost:8020", + "scopes": ["all-apis"], + "profile": "p", + }, + } + + assert copilot._token_provider(state)(False) == "custom-token" + assert calls == [(WS, "cid", "http://localhost:8020", ["all-apis"], "p", False)] + + +class _ExitedProcess: + def __init__(self, returncode=0, interrupts=0): + self.returncode = returncode + self.interrupts = interrupts + + def wait(self): + if self.interrupts: + self.interrupts -= 1 + raise KeyboardInterrupt + return self.returncode + + +def _stub_launch(monkeypatch, process=None, token="dbx-token"): + """Run launch against the real proxy with a stubbed Copilot process; return what it saw.""" + seen = {} + monkeypatch.setattr(copilot, "get_databricks_token", lambda *args, **kwargs: token) + monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") + + def popen(argv, *, env): + seen["argv"] = argv + seen["env"] = env + base = env["COPILOT_PROVIDER_BASE_URL"] + seen["accepted"] = httpx.post( + f"{base}/chat/completions", + headers={"Authorization": "Bearer wrong-secret"}, + content=b"{}", + timeout=5, + ).status_code + return process or _ExitedProcess() + + monkeypatch.setattr(copilot.subprocess_cross_os, "popen", popen) + return seen + class TestLaunch: @pytest.mark.parametrize( @@ -368,19 +557,7 @@ def test_launch_model_precedence_preserves_argv( ): isolate_copilot_config_paths(monkeypatch, tmp_path) monkeypatch.delenv("COPILOT_PROVIDER_WIRE_MODEL", raising=False) - monkeypatch.setattr(copilot, "get_databricks_token", lambda *args, **kwargs: "tok") - monkeypatch.setattr(copilot, "TOKEN_REFRESH_INTERVAL_SECONDS", 3600) - calls = [] - - class Process: - def wait(self): - return 0 - - def popen(argv, *, env): - calls.append((argv, env)) - return Process() - - monkeypatch.setattr(copilot.subprocess_cross_os, "popen", popen) + seen = _stub_launch(monkeypatch) state = {"workspace": WS, "copilot_default_model": default} with pytest.raises(SystemExit) as exit_info: @@ -391,52 +568,78 @@ def popen(argv, *, env): ) assert exit_info.value.code == 0 - argv, env = calls[0] - assert argv == ["copilot", *tool_args] - assert env["COPILOT_MODEL"] == expected_model - assert env["COPILOT_PROVIDER_WIRE_API"] == expected_api + assert seen["argv"] == ["copilot", *tool_args] + assert seen["env"]["COPILOT_MODEL"] == expected_model + assert seen["env"]["COPILOT_PROVIDER_WIRE_API"] == expected_api assert state["copilot_default_model"] == default - def test_token_refresh_keeps_launch_model(self, tmp_path, monkeypatch): - env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) + def test_copilot_gets_the_proxy_and_a_session_secret_not_the_token(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) monkeypatch.delenv("COPILOT_PROVIDER_WIRE_MODEL", raising=False) - monkeypatch.setattr(copilot, "TOKEN_REFRESH_INTERVAL_SECONDS", 0.01) - token_calls = [] - refreshed = threading.Event() + seen = _stub_launch(monkeypatch, token="dbx-token") - def get_token(*args, force_refresh=False, **kwargs): - token_calls.append(force_refresh) - if force_refresh: - refreshed.set() - return "tok" + with pytest.raises(SystemExit): + copilot.launch( + {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") + ) - monkeypatch.setattr(copilot, "get_databricks_token", get_token) + env = seen["env"] + assert env["COPILOT_PROVIDER_BASE_URL"].startswith("http://127.0.0.1:") + assert env["COPILOT_PROVIDER_BEARER_TOKEN"] not in ("", "dbx-token") + assert "dbx-token" not in env.values() + # The proxy is live while Copilot runs and refuses other credentials. + assert seen["accepted"] == 401 - class Process: - def wait(self): - assert refreshed.wait(timeout=2) - return 0 + def test_launch_persists_no_state_and_no_token(self, tmp_path, monkeypatch): + env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) + _stub_launch(monkeypatch, token="dbx-token") + saved = [] + monkeypatch.setattr(copilot, "save_state", saved.append) - monkeypatch.setattr( - copilot.subprocess_cross_os, - "popen", - lambda *args, **kwargs: Process(), - ) + with pytest.raises(SystemExit): + copilot.launch( + {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") + ) + + assert saved == [] + assert not env_path.exists() + + def test_ctrl_c_keeps_waiting_for_copilot(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + _stub_launch(monkeypatch, process=_ExitedProcess(returncode=130, interrupts=2)) with pytest.raises(SystemExit) as exit_info: copilot.launch( - {"workspace": WS, "copilot_default_model": "gpt-5"}, - [], - options=copilot.LaunchOptions(user_pinned_model="gpt-6-astra"), + {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") ) - written = copilot.parse_dotenv(env_path) - assert exit_info.value.code == 0 - assert token_calls[0] is False - assert token_calls[1:] - assert all(token_calls[1:]) - assert written["COPILOT_MODEL"] == "gpt-6-astra" - assert written["COPILOT_PROVIDER_WIRE_API"] == "responses" + assert exit_info.value.code == 130 + + def test_proxy_stops_when_copilot_cannot_start(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + monkeypatch.setattr(copilot, "get_databricks_token", lambda *args, **kwargs: "tok") + monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") + seen = {} + + def popen(argv, *, env): + seen["base"] = env["COPILOT_PROVIDER_BASE_URL"] + raise FileNotFoundError("copilot") + + monkeypatch.setattr(copilot.subprocess_cross_os, "popen", popen) + + with pytest.raises(FileNotFoundError): + copilot.launch( + {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") + ) + + with pytest.raises(httpx.ConnectError): + httpx.post(f"{seen['base']}/responses", content=b"{}", timeout=2) + + def test_no_model_is_a_launch_error(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + + with pytest.raises(RuntimeError, match="No Copilot model"): + copilot.launch({"workspace": WS}, [], options=copilot.LaunchOptions()) class TestManagedKeys: @@ -446,12 +649,14 @@ def test_includes_required_vars(self): "COPILOT_PROVIDER_BASE_URL", "COPILOT_PROVIDER_WIRE_API", "COPILOT_MODEL", - "COPILOT_PROVIDER_BEARER_TOKEN", "COPILOT_OFFLINE", - "OAUTH_TOKEN", ): assert key in copilot.MANAGED_KEYS + def test_excludes_credentials(self): + for key in copilot.CREDENTIAL_ENV_KEYS: + assert key not in copilot.MANAGED_KEYS + class TestValidateCmd: def test_starts_with_binary(self): diff --git a/tests/test_e2e.py b/tests/test_e2e.py index adc406947..da89d37f6 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -13,10 +13,12 @@ import atexit import json import os +import secrets import shutil import subprocess import sys import tempfile +import threading from pathlib import Path from urllib import error as urllib_error from urllib import request as urllib_request @@ -1098,27 +1100,41 @@ def test_launch_copilot_per_model( monkeypatch.setattr(copilot, "COPILOT_ENV_PATH", env_path) monkeypatch.setattr(copilot, "COPILOT_BACKUP_PATH", backup_path) - failures = [] - for family, model in models: - with pytest.MonkeyPatch().context() as mp: - mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.copilot.get_databricks_token", - lambda ws, profile=None, **kwargs: e2e_token, - ) - copilot.write_tool_config( - {**e2e_state, "workspace": e2e_workspace}, model, token=e2e_token - ) + from ucode import gateway_proxy + + # Copilot only ever sees the loopback proxy and a session secret; the proxy + # forwards to the real gateway with the e2e token. + session_secret = secrets.token_urlsafe(32) + server, cache, client = gateway_proxy.start_gateway_auth_proxy( + e2e_workspace, + lambda force_refresh: e2e_token, + session_secret=session_secret, + added_headers=copilot.proxy_headers({}), + sign_in_hint="Check the e2e token.", + ) + threading.Thread(target=server.serve_forever, daemon=True).start() + proxy_base_url = f"http://127.0.0.1:{server.server_address[1]}" - env = copilot.build_runtime_env(e2e_workspace, model, e2e_token) - cmd = copilot.validate_cmd("copilot") - result = _run_agent(cmd, env=env, timeout=120) - combined = (result.stdout + result.stderr).strip() - if result.returncode != 0 or not combined: - failures.append( - f"family={family} model={model} rc={result.returncode} " - f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" - ) + failures = [] + try: + for family, model in models: + with pytest.MonkeyPatch().context() as mp: + mp.setattr("ucode.agents.copilot.save_state", lambda s: None) + copilot.write_tool_config({**e2e_state, "workspace": e2e_workspace}, model) + + env = copilot.build_runtime_env(proxy_base_url, model, session_secret) + cmd = copilot.validate_cmd("copilot") + result = _run_agent(cmd, env=env, timeout=120) + combined = (result.stdout + result.stderr).strip() + if result.returncode != 0 or not combined: + failures.append( + f"family={family} model={model} rc={result.returncode} " + f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" + ) + finally: + cache.stop() + server.shutdown() + client.close() assert not failures, "Copilot launch failures:\n" + "\n".join(failures) diff --git a/tests/test_gateway_proxy_integration.py b/tests/test_gateway_proxy_integration.py index e1e57aece..3518f9865 100644 --- a/tests/test_gateway_proxy_integration.py +++ b/tests/test_gateway_proxy_integration.py @@ -47,10 +47,16 @@ class _FakeGateway: a streaming relay is exercised, not just a single write).""" def __init__( - self, status: int = 200, chunks: list[bytes] | None = None, sse_delay: float = 0.0 + self, + status: int = 200, + chunks: list[bytes] | None = None, + sse_delay: float = 0.0, + statuses: list[int] | None = None, ): self.requests: list[_CapturedRequest] = [] self._status = status + # Per-request statuses, in order; `status` applies once they run out. + self._statuses = list(statuses or []) self._chunks = chunks if chunks is not None else [b'{"ok":true}'] self._sse_delay = sse_delay self._server: ThreadingHTTPServer | None = None @@ -64,6 +70,7 @@ def base_url(self) -> str: def start(self) -> None: captured = self.requests status, chunks, sse_delay = self._status, self._chunks, self._sse_delay + statuses = self._statuses class Handler(BaseHTTPRequestHandler): def _serve(self) -> None: @@ -77,7 +84,7 @@ def _serve(self) -> None: body=body, ) ) - self.send_response(status) + self.send_response(statuses.pop(0) if statuses else status) self.send_header("Content-Type", "text/event-stream") self.end_headers() for chunk in chunks: @@ -263,3 +270,179 @@ def test_otlp_body_forwards_with_token_in_authorization(self, make_gateway): assert req.header("Authorization") == "Bearer per-user-otel-token" assert req.header("X-Databricks-AI-Gateway-Token") is None assert req.body == b"\x0a\x02\x08\x01" + + +@contextlib.contextmanager +def _running_gateway_auth_proxy( + gateway: _FakeGateway, + token_provider=None, + *, + added_headers=None, + on_refresh_failure=None, +): + server, cache, client = gateway_proxy.start_gateway_auth_proxy( + gateway.base_url, + token_provider or _counting_token(), + session_secret=_SESSION_SECRET, + added_headers=added_headers or {}, + sign_in_hint="Run `databricks auth login`.", + on_refresh_failure=on_refresh_failure or (lambda exc: None), + ) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + cache.stop() + client.close() + thread.join(timeout=2) + + +_SESSION_SECRET = "session-secret" + + +class TestGatewayAuthProxyEndToEnd: + def test_swaps_session_secret_for_fresh_token_and_adds_headers(self, make_gateway): + gw = make_gateway() + with _running_gateway_auth_proxy( + gw, + _counting_token("dbx-token"), + added_headers={"User-Agent": "ucode/1 copilot/2", "X-Tag": "t"}, + ) as proxy_url: + resp = httpx.post( + f"{proxy_url}/ai-gateway/mlflow/v1/chat/completions", + headers={ + "Authorization": f"Bearer {_SESSION_SECRET}", + "User-Agent": "OpenAI/JS 5.20.1", + "x-tag": "client-value", + }, + content=b'{"model":"m"}', + timeout=10, + ) + assert resp.status_code == 200 + req = gw.requests[-1] + assert req.path == "/ai-gateway/mlflow/v1/chat/completions" + assert req.header("Authorization") == "Bearer dbx-token" + assert req.header("User-Agent") == "ucode/1 copilot/2" + assert req.header("X-Tag") == "t" + assert req.body == b'{"model":"m"}' + + def test_accepts_secret_in_x_api_key_and_does_not_forward_it(self, make_gateway): + gw = make_gateway() + with _running_gateway_auth_proxy(gw, _counting_token("dbx-token")) as proxy_url: + resp = httpx.post( + f"{proxy_url}/ai-gateway/anthropic/v1/messages", + headers={"x-api-key": _SESSION_SECRET}, + content=b"{}", + timeout=10, + ) + assert resp.status_code == 200 + req = gw.requests[-1] + assert req.path == "/ai-gateway/anthropic/v1/messages" + assert req.header("x-api-key") is None + assert req.header("Authorization") == "Bearer dbx-token" + + @pytest.mark.parametrize("credential", [None, "Bearer wrong", "Basic session-secret"]) + def test_rejects_missing_or_wrong_secret_without_forwarding(self, make_gateway, credential): + gw = make_gateway() + headers = {"Authorization": credential} if credential else {} + with _running_gateway_auth_proxy(gw) as proxy_url: + resp = httpx.post( + f"{proxy_url}/ai-gateway/mlflow/v1/responses", + headers=headers, + content=b"{}", + timeout=10, + ) + assert resp.status_code == 401 + assert gw.requests == [] + + @pytest.mark.parametrize( + "path", + [ + "/api/2.0/clusters/list", + "/ai-gateway/anthropic/../../api/2.0/clusters/list", + "/ai-gateway/mlflow/v1/%2e%2e/%2e%2e/%2e%2e/api/2.0/clusters/list", + "/ai-gateway/mlflow/v1/%252e%252e/%252e%252e/%252e%252e/api/2.0/clusters/list", + "/ai-gateway/otel/v1/traces", + ], + ) + def test_refuses_paths_outside_model_routes(self, make_gateway, path): + gw = make_gateway() + with _running_gateway_auth_proxy(gw) as proxy_url: + resp = httpx.post( + f"{proxy_url}{path}", + headers={"Authorization": f"Bearer {_SESSION_SECRET}"}, + content=b"{}", + timeout=10, + ) + assert resp.status_code == 404 + assert gw.requests == [] + + def test_expired_token_is_refreshed_and_request_retried(self, make_gateway): + # The session outlives the token: the gateway rejects the first one, the proxy + # mints a new one and the same request succeeds without the agent restarting. + tokens = iter(["expired-token", "fresh-token"]) + calls: list[bool] = [] + + def provider(force_refresh: bool) -> str: + calls.append(force_refresh) + return next(tokens) + + gw = make_gateway(statuses=[401, 200]) + with _running_gateway_auth_proxy(gw, provider) as proxy_url: + resp = httpx.post( + f"{proxy_url}/ai-gateway/anthropic/v1/messages", + headers={"Authorization": f"Bearer {_SESSION_SECRET}"}, + content=b"{}", + timeout=10, + ) + assert resp.status_code == 200 + assert calls == [True, True] + assert [req.header("Authorization") for req in gw.requests] == [ + "Bearer expired-token", + "Bearer fresh-token", + ] + + @pytest.mark.parametrize( + ("path", "expected_error"), + [ + ( + "/ai-gateway/anthropic/v1/messages", + {"type": "error", "error": {"type": "authentication_error"}}, + ), + ("/ai-gateway/mlflow/v1/responses", {"error": {"type": "invalid_request_error"}}), + ], + ) + def test_failed_refresh_returns_actionable_401( + self, make_gateway, capsys, path, expected_error + ): + minted: list[bool] = [] + + def provider(force_refresh: bool) -> str: + minted.append(force_refresh) + if len(minted) > 1: + raise RuntimeError("refresh token revoked") + return "token" + + failures: list[str] = [] + gw = make_gateway(status=401) + with _running_gateway_auth_proxy( + gw, provider, on_refresh_failure=lambda exc: failures.append(str(exc)) + ) as proxy_url: + resp = httpx.post( + f"{proxy_url}{path}", + headers={"Authorization": f"Bearer {_SESSION_SECRET}"}, + content=b"{}", + timeout=10, + ) + assert resp.status_code == 401 + body = resp.json() + error = body["error"] + assert error["type"] == expected_error["error"]["type"] + assert body.get("type") == expected_error.get("type") + assert "Databricks sign-in expired. Run `databricks auth login`." in error["message"] + assert failures == ["refresh token revoked"] + assert len(gw.requests) == 1 + assert capsys.readouterr().err == "" diff --git a/tests/test_integration_contract.py b/tests/test_integration_contract.py index d9fe6f15f..9a829c470 100644 --- a/tests/test_integration_contract.py +++ b/tests/test_integration_contract.py @@ -269,6 +269,7 @@ def test_cuj_discovery_fails_for_a_file_without_a_check_name(tmp_path): "headless-windows", "full", "opencode", + "copilot", "managed", "managed-opencode", "dedicated-cuj-plan", @@ -591,7 +592,7 @@ def test_live_integration_cases_belong_to_exactly_one_ci_agent(): if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"): marks = module_marks | _markers(node.decorator_list) if marks & {"live", "managed_fixture", "workspace_switch"}: - assert len(marks & {"claude", "codex", "opencode"}) == 1, node.name + assert len(marks & {"claude", "codex", "opencode", "copilot"}) == 1, node.name def test_model_discovery_cases_match_current_launch_contract(): diff --git a/tests/test_integration_runner.py b/tests/test_integration_runner.py index 9730fb83d..87d555c08 100644 --- a/tests/test_integration_runner.py +++ b/tests/test_integration_runner.py @@ -31,6 +31,26 @@ def test_headless_only_selects_exact_nodes_for_requested_agents(tmp_path): ] +def test_copilot_headless_target_and_exact_version(tmp_path): + assert runner.integration_test_targets( + tmp_path, + ["copilot"], + platform_name="posix", + installation_only=False, + headless_only=True, + ) == [ + f"{tmp_path / 'test_ug_copilot_headless.py'}::test_ug_copilot_headless_completes_file_task" + ] + args = runner.arguments(["--copilot-version", "1.0.91", "--installation-only"], environment={}) + assert args.copilot_version == "1.0.91" + + +def test_copilot_version_output_matches_the_pinned_version(): + output = "GitHub Copilot CLI 1.0.91.\nRun 'copilot update' to check for updates.\n" + assert runner.agent_version_output("copilot", output) == "GitHub Copilot CLI 1.0.91" + assert runner.agent_version_output("codex", "codex-cli 0.154.0") == "codex-cli 0.154.0" + + def test_windows_live_runs_select_every_module_except_unsupported_ones(tmp_path): (tmp_path / "test_headless.py").write_text("import json\n") (tmp_path / "test_tui.py").write_text("from utils.terminal import AgentTerminal\n") From 204cc8d5aa3efab8858e57299991e8f5b4a53a60 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Mon, 19 Oct 2026 21:02:32 +0000 Subject: [PATCH 02/11] Close the Copilot proxy socket on exit and allow for Windows' slow refused connect Co-authored-by: Isaac --- src/ucode/agents/copilot.py | 1 + tests/test_agent_copilot.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 1984d4811..048b2715c 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -357,6 +357,7 @@ def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None finally: cache.stop() server.shutdown() + server.server_close() client.close() raise SystemExit(returncode) diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 97359c0dd..a983cd612 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -632,8 +632,9 @@ def popen(argv, *, env): {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") ) + # Windows retries a refused loopback connect for about 2s before failing. with pytest.raises(httpx.ConnectError): - httpx.post(f"{seen['base']}/responses", content=b"{}", timeout=2) + httpx.post(f"{seen['base']}/responses", content=b"{}", timeout=10) def test_no_model_is_a_launch_error(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) From 1532e9ff6f84446c34eb6c105407a4fae742c4d3 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Mon, 19 Oct 2026 21:02:59 +0000 Subject: [PATCH 03/11] Drop test comment Co-authored-by: Isaac --- tests/test_agent_copilot.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index a983cd612..568e7e472 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -632,7 +632,6 @@ def popen(argv, *, env): {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") ) - # Windows retries a refused loopback connect for about 2s before failing. with pytest.raises(httpx.ConnectError): httpx.post(f"{seen['base']}/responses", content=b"{}", timeout=10) From 6773cf6f46594f03b13b62bed35e292e5df66fe0 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Mon, 19 Oct 2026 21:47:56 +0000 Subject: [PATCH 04/11] copilot: integration journeys for token refresh, expired sign-in, and revert Co-authored-by: Isaac --- tests/README.md | 2 + tests/integration/README.md | 12 +- tests/integration/test_ug_copilot_headless.py | 30 +-- .../integration/test_ug_copilot_lifecycle.py | 195 ++++++++++++++++++ tests/integration/utils/constants.py | 8 + tests/integration/utils/evidence.py | 16 ++ tests/integration/utils/pat_bearer.py | 66 ++++++ 7 files changed, 299 insertions(+), 30 deletions(-) create mode 100644 tests/integration/test_ug_copilot_lifecycle.py create mode 100644 tests/integration/utils/pat_bearer.py diff --git a/tests/README.md b/tests/README.md index b01dd5cea..da91f4118 100644 --- a/tests/README.md +++ b/tests/README.md @@ -433,6 +433,8 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_ug_claude_headless_prompt_argument`, `test_ug_claude_headless_prompt_stdin`, `test_ug_claude_headless_prompt_after_separator` | Run Claude from a script using each prompt form | Structured final answer contains the file value; exit zero; no routing | | `test_ug_codex_headless_prompt_argument`, `test_ug_codex_headless_prompt_stdin`, `test_ug_codex_headless_prompt_after_separator` | Run Codex from a script using each prompt form | Completed turn and final answer contain the file value; exit zero; no routing | | `test_ug_copilot_headless_completes_file_task`, `test_ug_copilot_headless_gpt_responses_task` | Run Copilot CLI from a script (`-p ... --output-format json`) on a Claude model (Anthropic Messages) and a GPT model (Responses) | Completed turn whose assistant message contains the file value; no file under the session home holds the Databricks credential (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. Tokens written by older ug versions: unit tests only (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 | | `test_ug_agents_admin_managed_guardrails` | `ug agents add` / `remove` on an agent the admin config enables | Add is a no-op noting the admin manages it; remove is rejected with nonzero exit | diff --git a/tests/integration/README.md b/tests/integration/README.md index fada250fd..fe5fd4d1a 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -206,8 +206,11 @@ model-selection journey. Copilot CLI has two live headless journeys in `test_ug_copilot_headless.py`, one per wire API (Anthropic Messages for Claude, Responses for GPT). Both go through the session's loopback proxy and check that no file under the session home holds the -Databricks credential. Token refresh past expiry is covered by the hermetic proxy -tests in `tests/test_gateway_proxy_integration.py`, not by a live hour-long session. +Databricks credential. `test_ug_copilot_lifecycle.py` signs in through +`DATABRICKS_BEARER_COMMAND` (`utils/pat_bearer.py` mints 15-minute PATs) and has the agent +revoke them mid-task: once to require a refresh, once with sign-out to require the +sign-in-expired message. It also covers configure/revert. Refresh near natural expiry is +covered by the hermetic proxy tests in `tests/test_gateway_proxy_integration.py`. ## Shared agent helpers @@ -422,6 +425,7 @@ test_ug_claude_relayed.py # relayed session: subscription + Databr test_ug_claude_tracing.py # Claude OTLP export reaches the configured trace table test_ug_codex_headless.py # script prompts and model arguments test_ug_copilot_headless.py # Copilot via the session proxy: Claude + GPT tasks, no token on disk +test_ug_copilot_lifecycle.py # Copilot refresh after revocation, sign-in expired, configure/revert test_ug_codex_tracing.py # Codex OTLP export reaches the configured trace table test_ug_claude_commands.py # command help forwarding test_ug_codex_commands.py # command help and parser error forwarding @@ -615,7 +619,7 @@ cases instead complete real file tasks through `--workspace` and `--model-locati journeys check setup and launch with existing services, without inference. There are **85 live cases** (including 11 marked TUI journeys) and **7 installation -checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds two live headless cases. One +checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds five live cases. One **`workspace_switch` case** uses two real workspaces and checks skills MCP cleanup and a completed Claude task. A further **43 `managed_fixture` cases** (ten of them also `live`) run on the managed workspace with a checked-in JSON CodingAgentConfig from `tests/fixtures/managed_config/` @@ -841,7 +845,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`, two cases) runs on Linux and Windows the +A non-blocking **Copilot** job (`live and copilot`, five cases) runs on Linux and Windows the same way and is likewise outside `cujs` until it is stable. Each lane installs only the agent CLIs it needs (one, except the managed OpenCode lane's diff --git a/tests/integration/test_ug_copilot_headless.py b/tests/integration/test_ug_copilot_headless.py index 1527d8f19..cae6f08a7 100644 --- a/tests/integration/test_ug_copilot_headless.py +++ b/tests/integration/test_ug_copilot_headless.py @@ -1,33 +1,11 @@ """CUJs for using GitHub Copilot CLI from scripts through installed ug.""" import pytest -from utils.constants import CLAUDE_TEST_MODEL, CODEX_TEST_MODEL -from utils.evidence import FileTask +from utils.constants import CLAUDE_TEST_MODEL, CODEX_TEST_MODEL, COPILOT_HEADLESS_FLAGS +from utils.evidence import FileTask, assert_no_bearer_on_disk pytestmark = [pytest.mark.live, pytest.mark.copilot] -COPILOT_HEADLESS_FLAGS = ( - "--allow-all-tools", - "--output-format", - "json", - "--no-auto-update", - "--disable-builtin-mcps", -) - - -def _assert_no_credential_on_disk(session): - """No Databricks credential in any file ug or Copilot wrote under the session home.""" - bearer = session.env.get("DATABRICKS_BEARER") - assert bearer, "The live suite authenticates with DATABRICKS_BEARER." - for root in (session.home / ".copilot", session.home / ".ucode"): - for path in root.rglob("*") if root.exists() else (): - if path.is_file(): - text = path.read_text(encoding="utf-8", errors="replace") - assert bearer not in text, f"{path} holds the Databricks credential" - env_file = (session.home / ".copilot/ucode.env").read_text(encoding="utf-8") - assert "COPILOT_PROVIDER_BEARER_TOKEN" not in env_file - assert "OAUTH_TOKEN" not in env_file - def test_ug_copilot_headless_completes_file_task(live_session, workspace): """Scenario: configure Copilot, then run a headless Claude prompt through `ug copilot`. @@ -61,7 +39,7 @@ def test_ug_copilot_headless_completes_file_task(live_session, workspace): ) task.assert_headless_answer("copilot", result) - _assert_no_credential_on_disk(session) + assert_no_bearer_on_disk(session) def test_ug_copilot_headless_gpt_responses_task(live_session, workspace): @@ -96,4 +74,4 @@ def test_ug_copilot_headless_gpt_responses_task(live_session, workspace): ) task.assert_headless_answer("copilot", result) - _assert_no_credential_on_disk(session) + assert_no_bearer_on_disk(session) diff --git a/tests/integration/test_ug_copilot_lifecycle.py b/tests/integration/test_ug_copilot_lifecycle.py new file mode 100644 index 000000000..02198be20 --- /dev/null +++ b/tests/integration/test_ug_copilot_lifecycle.py @@ -0,0 +1,195 @@ +"""CUJs for Copilot's ug sign-in over a session, and for undoing its ug setup.""" + +import json +import os +import shlex +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +from utils.constants import CLAUDE_TEST_MODEL, CODEX_TEST_MODEL, COPILOT_HEADLESS_FLAGS +from utils.evidence import FileTask, assert_no_bearer_on_disk, fresh_id +from utils.terminal import TerminalProcess + +pytestmark = [pytest.mark.live, pytest.mark.copilot] + +PAT_BEARER = Path(__file__).parent / "utils" / "pat_bearer.py" + + +def _configure(session, workspace): + session.run( + "configure", + "--agents", + "copilot", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + + +def _pat_sign_in(session, workspace): + """Switch the session to ug's bearer-command sign-in, minting short-lived PATs. + + Returns the state dir and the shell step Copilot runs to revoke or sign out; the step + uses only paths under the session's working directory.""" + state = session.cwd / f"auth-{fresh_id()}" + state.mkdir() + shutil.copy(PAT_BEARER, state / PAT_BEARER.name) + argv = [sys.executable, "-I", str(state / PAT_BEARER.name), "mint", workspace, str(state)] + command = subprocess.list2cmdline(argv) if os.name == "nt" else shlex.join(argv) + session.env["UG_INTEGRATION_ISSUER_BEARER"] = session.env.pop("DATABRICKS_BEARER") + session.env["DATABRICKS_BEARER_COMMAND"] = command + python = "python" if os.name == "nt" else "python3" + relative = f"{state.name}/{PAT_BEARER.name}" + + def step(mode): + return f"{python} {relative} {mode} {workspace} {state.name}" + + return state, step + + +def _delete_minted(session, workspace, state): + subprocess.run( + [sys.executable, "-I", str(state / PAT_BEARER.name), "cleanup", workspace, str(state)], + env=session.env, + stdin=subprocess.DEVNULL, + capture_output=True, + timeout=120, + check=False, + ) + + +def _minted_count(state): + return len((state / "minted").read_text(encoding="utf-8").split()) + + +def test_ug_copilot_refreshes_a_revoked_token_mid_session(live_session, workspace): + """Scenario: sign in through a bearer command, then revoke the session's token mid-task. + + The headless Claude task first runs a shell step that deletes every PAT minted so far, + including the one ug's proxy holds, then reads the fixture. + + Expected: the gateway rejects the revoked token, ug mints a new one without restarting + Copilot, and Copilot's completed answer contains the fixture's value. Limits: needs the + e2e principal to create PATs; a gateway that still honors a just-deleted PAT fails the + mint-after-revoke assertion. + """ + session = live_session + task = FileTask(session) + _configure(session, workspace) + state, step = _pat_sign_in(session, workspace) + session.record("copilot-model.json", {"model": CLAUDE_TEST_MODEL}) + prompt = f"First run this shell command exactly once: {step('revoke')} Then: {task.prompt}" + + try: + result = session.run( + "copilot", + "--model", + CLAUDE_TEST_MODEL, + "-p", + prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=240, + ) + + task.assert_headless_answer("copilot", result) + revoked = int((state / "revoked").read_text(encoding="utf-8")) + session.record( + "copilot-token-mints.json", {"revoked": revoked, "minted": _minted_count(state)} + ) + assert _minted_count(state) > revoked, "No token was minted after the revocation" + finally: + _delete_minted(session, workspace, state) + + +def test_ug_copilot_reports_expired_sign_in(live_session, workspace): + """Scenario: sign in through a bearer command, then sign out and revoke mid-task. + + The headless Claude task first runs a shell step that makes the bearer command fail + and deletes every PAT minted so far, then reads the fixture. + + Expected: ug cannot mint a replacement, so Copilot's output shows ug's sign-in-expired + message with the `databricks auth login` fix, and the failure is logged to + ~/.ucode/logs/copilot-proxy.log instead of Copilot's terminal. + """ + session = live_session + task = FileTask(session) + _configure(session, workspace) + state, step = _pat_sign_in(session, workspace) + session.record("copilot-model.json", {"model": CLAUDE_TEST_MODEL}) + prompt = f"First run this shell command exactly once: {step('sign-out')} Then: {task.prompt}" + + try: + result = session.run( + "copilot", + "--model", + CLAUDE_TEST_MODEL, + "-p", + prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=240, + ok=False, + ) + + assert (state / "revoked").exists(), "Copilot never ran the sign-out step" + output = result.stdout + result.stderr + assert "Databricks sign-in expired" in output, output + assert "databricks auth login" in output, output + log = (session.home / ".ucode/logs/copilot-proxy.log").read_text(encoding="utf-8") + assert "token refresh failed" in log + assert "token refresh failed" not in output + finally: + _delete_minted(session, workspace, state) + + +def test_ug_configure_copilot_repeat_and_revert(live_session, workspace): + """Scenario: configure Copilot twice over a user-owned MCP config, run a task, then revert. + + Expected: repeat setup stays usable, a headless GPT task completes, no file under the + session home holds the bearer, revert removes ug's generated Copilot env file and leaves + the user's file byte-identical, and a repeated revert succeeds. Limits: tokens written + by older ug versions are scrubbed by `configure`/`revert`, but this suite has no older ug + to produce them; unit tests cover that path. + """ + session = live_session + task = FileTask(session) + # Ordinary pre-existing user settings, not manufactured gateway state. + user_path = session.home / ".copilot/mcp-config.json" + original = json.dumps({"mcpServers": {}}, indent=2) + "\n" + user_path.parent.mkdir(parents=True) + user_path.write_text(original, encoding="utf-8") + + _configure(session, workspace) + first = session.workspace_state() + assert "copilot" in first["available_tools"] + _configure(session, workspace) + assert session.workspace_state()["available_tools"] == first["available_tools"] + assert_no_bearer_on_disk(session) + session.record("copilot-model.json", {"model": CODEX_TEST_MODEL}) + + result = session.run( + "copilot", + "--model", + CODEX_TEST_MODEL, + "-p", + task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + task.assert_headless_answer("copilot", result) + assert_no_bearer_on_disk(session) + + with TerminalProcess(session, "ug", [str(session.binary), "revert"], "revert") as terminal: + terminal.finish() + assert user_path.read_text(encoding="utf-8") == original + assert not (session.home / ".copilot/ucode.env").exists() + assert "Not Configured" in session.run("status").stdout + assert_no_bearer_on_disk(session) + with TerminalProcess( + session, "ug", [str(session.binary), "revert"], "repeat-revert" + ) as terminal: + terminal.finish() diff --git a/tests/integration/utils/constants.py b/tests/integration/utils/constants.py index 1783073db..a3adbd04c 100644 --- a/tests/integration/utils/constants.py +++ b/tests/integration/utils/constants.py @@ -6,5 +6,13 @@ # nano sometimes nests powershell inside Codex's pwsh, breaking `$` expansion in its file reads. CODEX_TEST_MODEL = "system.ai.gpt-5-6-luna" if os.name == "nt" else "system.ai.gpt-5-4-nano" +COPILOT_HEADLESS_FLAGS = ( + "--allow-all-tools", + "--output-format", + "json", + "--no-auto-update", + "--disable-builtin-mcps", +) + MANAGED_CLAUDE_PROVIDER_SERVICE = "main.default.ci_e2e_anthropic_mps" MANAGED_CODEX_PROVIDER_SERVICE = "main.default.ci_e2e_openai_mps" diff --git a/tests/integration/utils/evidence.py b/tests/integration/utils/evidence.py index f0fba4e38..bec0f6c61 100644 --- a/tests/integration/utils/evidence.py +++ b/tests/integration/utils/evidence.py @@ -132,6 +132,22 @@ def assistant_answer_contains(session, agent: str, value: str, *, child: bool = ) +def assert_no_bearer_on_disk(session) -> None: + """No file ug or Copilot wrote under the session home holds the Databricks bearer.""" + bearer = session.env.get("DATABRICKS_BEARER") + assert bearer, "The live suite authenticates with DATABRICKS_BEARER." + for root in (session.home / ".copilot", session.home / ".ucode"): + for path in root.rglob("*") if root.exists() else (): + if path.is_file(): + text = path.read_text(encoding="utf-8", errors="replace") + assert bearer not in text, f"{path} holds the Databricks credential" + env_file = session.home / ".copilot/ucode.env" + if env_file.exists(): + text = env_file.read_text(encoding="utf-8") + assert "COPILOT_PROVIDER_BEARER_TOKEN" not in text + assert "OAUTH_TOKEN" not in text + + def fresh_id(): """Return a 12-hex-char id; models drop characters when copying longer ones.""" return uuid.uuid4().hex[:12] diff --git a/tests/integration/utils/pat_bearer.py b/tests/integration/utils/pat_bearer.py new file mode 100644 index 000000000..4b4f544b4 --- /dev/null +++ b/tests/integration/utils/pat_bearer.py @@ -0,0 +1,66 @@ +"""Bearer command for the Copilot sign-in CUJs, run as ``pat_bearer.py MODE HOST STATE_DIR``. + +``mint`` prints a fresh short-lived workspace PAT (ug's ``DATABRICKS_BEARER_COMMAND`` contract) +issued with ``UG_INTEGRATION_ISSUER_BEARER``, unless ``sign-out`` has run. ``revoke`` and +``sign-out`` delete every PAT minted so far, as a user revoking or signing out mid-session would; +``cleanup`` deletes them at teardown. Token ids, never token values, are kept in STATE_DIR. +""" + +import json +import os +import subprocess +import sys +from pathlib import Path + + +def databricks(host: str, *args: str, check: bool = True) -> str: + env = {**os.environ, "DATABRICKS_HOST": host} + env["DATABRICKS_TOKEN"] = env.pop("UG_INTEGRATION_ISSUER_BEARER") + env.pop("DATABRICKS_BEARER", None) + env.pop("DATABRICKS_BEARER_COMMAND", None) + result = subprocess.run( + ["databricks", *args, "-o", "json"], + env=env, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=12, + check=False, + ) + if check and result.returncode: + sys.exit(f"databricks {args[0]} {args[1]} failed: {result.stderr.strip()[:300]}") + return result.stdout + + +def main() -> None: + mode, host, state = sys.argv[1], sys.argv[2], Path(sys.argv[3]) + minted = state / "minted" + if mode == "mint": + if (state / "signed-out").exists(): + sys.exit("signed out") + info = json.loads( + databricks( + host, + "tokens", + "create", + "--lifetime-seconds", + "900", + "--comment", + "ug-integration-copilot", + ) + ) + with minted.open("a", encoding="utf-8") as ids: + ids.write(info["token_info"]["token_id"] + "\n") + print(info["token_value"]) + return + if mode == "sign-out": + (state / "signed-out").touch() + token_ids = minted.read_text(encoding="utf-8").split() if minted.exists() else [] + for token_id in token_ids: + databricks(host, "tokens", "delete", token_id, check=mode != "cleanup") + if mode != "cleanup": + (state / "revoked").write_text(str(len(token_ids)), encoding="utf-8") + + +if __name__ == "__main__": + main() From f6ebc196551cf3cd8aa3f46b93f5e72d0b3b6169 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Mon, 19 Oct 2026 21:55:14 +0000 Subject: [PATCH 05/11] copilot: integration journey for upgrading from a ug that saved the token Co-authored-by: Isaac --- tests/README.md | 3 +- tests/integration/README.md | 8 +- .../integration/test_ug_copilot_lifecycle.py | 75 ++++++++++++++++++- 3 files changed, 78 insertions(+), 8 deletions(-) diff --git a/tests/README.md b/tests/README.md index da91f4118..ba51b2caa 100644 --- a/tests/README.md +++ b/tests/README.md @@ -434,7 +434,8 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_ug_codex_headless_prompt_argument`, `test_ug_codex_headless_prompt_stdin`, `test_ug_codex_headless_prompt_after_separator` | Run Codex from a script using each prompt form | Completed turn and final answer contain the file value; exit zero; no routing | | `test_ug_copilot_headless_completes_file_task`, `test_ug_copilot_headless_gpt_responses_task` | Run Copilot CLI from a script (`-p ... --output-format json`) on a Claude model (Anthropic Messages) and a GPT model (Responses) | Completed turn whose assistant message contains the file value; no file under the session home holds the Databricks credential (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. Tokens written by older ug versions: unit tests only (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_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 | | `test_ug_agents_admin_managed_guardrails` | `ug agents add` / `remove` on an agent the admin config enables | Add is a no-op noting the admin manages it; remove is rejected with nonzero exit | diff --git a/tests/integration/README.md b/tests/integration/README.md index fe5fd4d1a..07e46b35c 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -209,7 +209,7 @@ session's loopback proxy and check that no file under the session home holds the Databricks credential. `test_ug_copilot_lifecycle.py` signs in through `DATABRICKS_BEARER_COMMAND` (`utils/pat_bearer.py` mints 15-minute PATs) and has the agent revoke them mid-task: once to require a refresh, once with sign-out to require the -sign-in-expired message. It also covers configure/revert. Refresh near natural expiry is +sign-in-expired message. It also covers configure/revert, including upgrading from ug v0.2.0, which saved the bearer. Refresh near natural expiry is covered by the hermetic proxy tests in `tests/test_gateway_proxy_integration.py`. ## Shared agent helpers @@ -425,7 +425,7 @@ test_ug_claude_relayed.py # relayed session: subscription + Databr test_ug_claude_tracing.py # Claude OTLP export reaches the configured trace table test_ug_codex_headless.py # script prompts and model arguments test_ug_copilot_headless.py # Copilot via the session proxy: Claude + GPT tasks, no token on disk -test_ug_copilot_lifecycle.py # Copilot refresh after revocation, sign-in expired, configure/revert +test_ug_copilot_lifecycle.py # Copilot refresh after revocation, sign-in expired, configure/revert, upgrade test_ug_codex_tracing.py # Codex OTLP export reaches the configured trace table test_ug_claude_commands.py # command help forwarding test_ug_codex_commands.py # command help and parser error forwarding @@ -619,7 +619,7 @@ cases instead complete real file tasks through `--workspace` and `--model-locati journeys check setup and launch with existing services, without inference. There are **85 live cases** (including 11 marked TUI journeys) and **7 installation -checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds five live cases. One +checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds six live cases. One **`workspace_switch` case** uses two real workspaces and checks skills MCP cleanup and a completed Claude task. A further **43 `managed_fixture` cases** (ten of them also `live`) run on the managed workspace with a checked-in JSON CodingAgentConfig from `tests/fixtures/managed_config/` @@ -845,7 +845,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`, five cases) runs on Linux and Windows the +A non-blocking **Copilot** job (`live and copilot`, six cases) runs on Linux and Windows the same way and is likewise outside `cujs` until it is stable. Each lane installs only the agent CLIs it needs (one, except the managed OpenCode lane's diff --git a/tests/integration/test_ug_copilot_lifecycle.py b/tests/integration/test_ug_copilot_lifecycle.py index 02198be20..6278b17b0 100644 --- a/tests/integration/test_ug_copilot_lifecycle.py +++ b/tests/integration/test_ug_copilot_lifecycle.py @@ -16,6 +16,11 @@ pytestmark = [pytest.mark.live, pytest.mark.copilot] PAT_BEARER = Path(__file__).parent / "utils" / "pat_bearer.py" +# v0.2.0, the last release that saved the Databricks token in ~/.copilot/ucode.env. +TOKEN_SAVING_UG = ( + "unity-gateway @ git+https://github.com/databricks/unity-gateway" + "@cb7bd8c54800ac59b9701ea0ac9f561c83d07bf7" +) def _configure(session, workspace): @@ -63,6 +68,25 @@ def _delete_minted(session, workspace, state): ) +def _install_token_saving_ug(session) -> Path: + uv = shutil.which("uv", path=session.env["PATH"]) + assert uv, "Installing the older ug release needs uv." + venv = session.home.parent / "ug-0.2.0" + for command in ( + [uv, "venv", "--python", "3.12", str(venv)], + [uv, "pip", "install", "--python", str(venv), TOKEN_SAVING_UG], + ): + subprocess.run( + command, + env=session.env, + stdin=subprocess.DEVNULL, + capture_output=True, + timeout=300, + check=True, + ) + return venv / ("Scripts/ug.exe" if os.name == "nt" else "bin/ug") + + def _minted_count(state): return len((state / "minted").read_text(encoding="utf-8").split()) @@ -151,9 +175,7 @@ def test_ug_configure_copilot_repeat_and_revert(live_session, workspace): Expected: repeat setup stays usable, a headless GPT task completes, no file under the session home holds the bearer, revert removes ug's generated Copilot env file and leaves - the user's file byte-identical, and a repeated revert succeeds. Limits: tokens written - by older ug versions are scrubbed by `configure`/`revert`, but this suite has no older ug - to produce them; unit tests cover that path. + the user's file byte-identical, and a repeated revert succeeds. """ session = live_session task = FileTask(session) @@ -193,3 +215,50 @@ def test_ug_configure_copilot_repeat_and_revert(live_session, workspace): session, "ug", [str(session.binary), "revert"], "repeat-revert" ) as terminal: terminal.finish() + + +def test_ug_copilot_upgrade_scrubs_token_saved_by_older_ug(live_session, workspace): + """Scenario: configure Copilot with ug v0.2.0, upgrade to this ug, configure, use, revert. + + Expected: v0.2.0 leaves the bearer in ~/.copilot/ucode.env (the scenario's input); + this ug's configure removes it so no file under the session home holds the bearer, + a headless GPT task completes, and revert leaves no token and no generated env file. + """ + session = live_session + task = FileTask(session) + token_saving_ug = _install_token_saving_ug(session) + env_file = session.home / ".copilot/ucode.env" + + session.run( + "configure", + "--agents", + "copilot", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + binary=token_saving_ug, + ) + assert session.env["DATABRICKS_BEARER"] in env_file.read_text(encoding="utf-8"), ( + "ug v0.2.0 no longer saves the bearer, so this journey has no legacy token to scrub" + ) + + _configure(session, workspace) + assert_no_bearer_on_disk(session) + session.record("copilot-model.json", {"model": CODEX_TEST_MODEL}) + result = session.run( + "copilot", + "--model", + CODEX_TEST_MODEL, + "-p", + task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + task.assert_headless_answer("copilot", result) + + with TerminalProcess(session, "ug", [str(session.binary), "revert"], "revert") as terminal: + terminal.finish() + assert not env_file.exists() + assert "Not Configured" in session.run("status").stdout + assert_no_bearer_on_disk(session) From d480b811eca64c0851a29da17368b5485cbc99df Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Tue, 20 Oct 2026 10:00:57 +0000 Subject: [PATCH 06/11] Share the launch token selection; review fixes in Copilot tests Co-authored-by: Isaac --- src/ucode/agents/codex.py | 15 ++----------- src/ucode/agents/copilot.py | 21 ++----------------- src/ucode/custom_oauth.py | 16 ++++++++++++++ .../integration/test_ug_copilot_lifecycle.py | 1 + tests/integration/utils/evidence.py | 11 +++++----- tests/test_agent_codex.py | 20 +++++++++--------- tests/test_agent_copilot.py | 11 +++++----- tests/test_cli.py | 2 +- tests/test_codex_smart_routing_v2.py | 4 ++-- tests/test_e2e.py | 1 + 10 files changed, 46 insertions(+), 56 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 3408c9365..c726ed5db 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -50,7 +50,7 @@ CUSTOM_OAUTH_TIMEOUT_MS, CustomOAuthConfig, build_custom_auth_token_argv, - get_custom_client_token, + launch_token, ) from ucode.databricks import ( CodexCatalogSource, @@ -58,7 +58,6 @@ _fetch_codex_model_catalog, build_auth_token_argv, build_tool_base_url, - get_databricks_token, ug_binary, ) from ucode.launcher import exec_or_spawn @@ -1181,17 +1180,7 @@ def _launch_token(state: dict, workspace: str, *, force_refresh: bool = False) - """The token Codex authenticates with: a custom-OAuth client token when configured, else the CLI profile. The single auth-selection point — the OTLP proxy's token provider (_otel_token_provider) delegates here so both mint the same principal.""" - custom_oauth = state.get("custom_oauth") - if isinstance(custom_oauth, dict): - return get_custom_client_token( - workspace, - custom_oauth["client_id"], - custom_oauth["redirect_url"], - scopes=custom_oauth["scopes"], - profile=custom_oauth.get("profile"), - force_refresh=force_refresh, - ) - return get_databricks_token(workspace, state.get("profile"), force_refresh=force_refresh) + return launch_token(state, workspace, force_refresh=force_refresh) def _tool_args_select_model(tool_args: list[str]) -> bool: diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 048b2715c..d08f7eba5 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -44,11 +44,10 @@ write_json_file, ) from ucode.constants import AI_GATEWAY_REQUEST_TAGS_HEADER, LOOPBACK_HOST -from ucode.custom_oauth import get_custom_client_token +from ucode.custom_oauth import launch_token from ucode.databricks import ( build_copilot_base_url, build_tool_base_url, - get_databricks_token, ) from ucode.os_compatibility import subprocess_cross_os from ucode.state import is_tool_managed, mark_tool_managed, save_state @@ -268,24 +267,8 @@ def write_tool_config(state: dict, model: str) -> dict: def _token_provider(state: dict) -> Callable[[bool], str]: - """Mint from the source the rest of ug uses for this workspace: custom OAuth when configured, - else the CLI profile.""" workspace = state["workspace"] - custom_oauth = state.get("custom_oauth") - - def provider(force_refresh: bool) -> str: - if isinstance(custom_oauth, dict): - return get_custom_client_token( - workspace, - custom_oauth["client_id"], - custom_oauth["redirect_url"], - scopes=custom_oauth["scopes"], - profile=custom_oauth.get("profile"), - force_refresh=force_refresh, - ) - return get_databricks_token(workspace, state.get("profile"), force_refresh=force_refresh) - - return provider + return lambda force_refresh: launch_token(state, workspace, force_refresh=force_refresh) def proxy_headers(state: dict) -> dict[str, str]: diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index 9a3d77931..8ba64667e 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -238,3 +238,19 @@ def get_custom_client_token( f"redirect URL ({config['redirect_url']}); ensure its local port is available and " "the SDK token cache is writable, then retry." ) from exc + + +def launch_token(state: dict, workspace: str, *, force_refresh: bool = False) -> str: + """The token an agent launch authenticates with: the custom-OAuth client token when + configured, else the CLI profile's.""" + custom_oauth = state.get("custom_oauth") + if isinstance(custom_oauth, dict): + return get_custom_client_token( + workspace, + custom_oauth["client_id"], + custom_oauth["redirect_url"], + scopes=custom_oauth["scopes"], + profile=custom_oauth.get("profile"), + force_refresh=force_refresh, + ) + return get_databricks_token(workspace, state.get("profile"), force_refresh=force_refresh) diff --git a/tests/integration/test_ug_copilot_lifecycle.py b/tests/integration/test_ug_copilot_lifecycle.py index 6278b17b0..42bd005bc 100644 --- a/tests/integration/test_ug_copilot_lifecycle.py +++ b/tests/integration/test_ug_copilot_lifecycle.py @@ -121,6 +121,7 @@ def test_ug_copilot_refreshes_a_revoked_token_mid_session(live_session, workspac ) task.assert_headless_answer("copilot", result) + assert (state / "revoked").exists(), "Copilot never ran the revoke step" revoked = int((state / "revoked").read_text(encoding="utf-8")) session.record( "copilot-token-mints.json", {"revoked": revoked, "minted": _minted_count(state)} diff --git a/tests/integration/utils/evidence.py b/tests/integration/utils/evidence.py index bec0f6c61..0e6e32c55 100644 --- a/tests/integration/utils/evidence.py +++ b/tests/integration/utils/evidence.py @@ -133,14 +133,13 @@ def assistant_answer_contains(session, agent: str, value: str, *, child: bool = def assert_no_bearer_on_disk(session) -> None: - """No file ug or Copilot wrote under the session home holds the Databricks bearer.""" + """No file under the session home holds the Databricks bearer.""" bearer = session.env.get("DATABRICKS_BEARER") assert bearer, "The live suite authenticates with DATABRICKS_BEARER." - for root in (session.home / ".copilot", session.home / ".ucode"): - for path in root.rglob("*") if root.exists() else (): - if path.is_file(): - text = path.read_text(encoding="utf-8", errors="replace") - assert bearer not in text, f"{path} holds the Databricks credential" + encoded = bearer.encode() + for path in session.home.rglob("*"): + if path.is_file() and not path.is_symlink(): + assert encoded not in path.read_bytes(), f"{path} holds the Databricks credential" env_file = session.home / ".copilot/ucode.env" if env_file.exists(): text = env_file.read_text(encoding="utf-8") diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index f887f4d20..8b35601ce 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -11,7 +11,7 @@ import pytest -from ucode import gateway_proxy, managed_files +from ucode import custom_oauth, gateway_proxy, managed_files from ucode.agents import LaunchOptions, codex from ucode.config_io import read_toml_safe from ucode.smart_routing import codex_routing @@ -996,7 +996,7 @@ def _patch(tmp_path, monkeypatch): monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) monkeypatch.setattr( - codex, + custom_oauth, "get_databricks_token", lambda workspace, profile=None, force_refresh=False: "tok", ) @@ -1010,7 +1010,7 @@ def test_sets_oauth_token(self, tmp_path, monkeypatch): monkeypatch.delenv("OAUTH_TOKEN", raising=False) launches = self._patch(tmp_path, monkeypatch) monkeypatch.setattr( - codex, + custom_oauth, "get_databricks_token", lambda workspace, profile=None, force_refresh=False: "fresh-token", ) @@ -1435,7 +1435,7 @@ def test_provider_discovery_uses_custom_oauth_token(self, tmp_path, monkeypatch) seen = {} monkeypatch.setattr(codex, "_model_catalog_path", lambda workspace, provider: catalog_path) monkeypatch.setattr( - codex, + custom_oauth, "get_databricks_token", lambda *args, **kwargs: pytest.fail("standard token used"), ) @@ -1446,7 +1446,7 @@ def custom_token( seen["profile"] = profile return "custom-token" - monkeypatch.setattr(codex, "get_custom_client_token", custom_token) + monkeypatch.setattr(custom_oauth, "get_custom_client_token", custom_token) def fetch(workspace, token, **kwargs): seen.update(workspace=workspace, token=token, provider=kwargs["identifier"]) @@ -1598,7 +1598,7 @@ def test_launches_legacy_config_after_configure( monkeypatch.setattr(codex, "LEGACY_CODEX_BACKUP_PATH", tmp_path / "legacy-backup.toml") monkeypatch.setattr(codex, "agent_version", lambda binary: version) monkeypatch.setattr(codex, "save_state", lambda state: None) - monkeypatch.setattr(codex, "get_databricks_token", lambda *_args, **_kw: "tok") + monkeypatch.setattr(custom_oauth, "get_databricks_token", lambda *_args, **_kw: "tok") monkeypatch.delenv("OAUTH_TOKEN", raising=False) launches: list[list[str]] = [] monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) @@ -1626,7 +1626,7 @@ def test_requires_populated_ucode_profile(self, tmp_path, monkeypatch): monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") launches = [] monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) - monkeypatch.setattr(codex, "get_databricks_token", lambda *_args, **_kw: "tok") + monkeypatch.setattr(custom_oauth, "get_databricks_token", lambda *_args, **_kw: "tok") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) with pytest.raises(RuntimeError, match="ug configure --agents codex"): @@ -2259,7 +2259,7 @@ def boom(path): class TestOtelTokenProvider: def test_token_provider_uses_default_profile(self, monkeypatch): get_token = Mock(return_value="token") - monkeypatch.setattr(codex, "get_databricks_token", get_token) + monkeypatch.setattr(custom_oauth, "get_databricks_token", get_token) provider = codex._otel_token_provider({"profile": "myprof"}, WS) assert provider(True) == "token" get_token.assert_called_once_with(WS, "myprof", force_refresh=True) @@ -2267,8 +2267,8 @@ def test_token_provider_uses_default_profile(self, monkeypatch): def test_token_provider_uses_custom_oauth_when_configured(self, monkeypatch): get_custom_token = Mock(return_value="custom-token") get_default_token = Mock() - monkeypatch.setattr(codex, "get_custom_client_token", get_custom_token) - monkeypatch.setattr(codex, "get_databricks_token", get_default_token) + monkeypatch.setattr(custom_oauth, "get_custom_client_token", get_custom_token) + monkeypatch.setattr(custom_oauth, "get_databricks_token", get_default_token) state = { "custom_oauth": { "client_id": "cid", diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 568e7e472..71e995ec4 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -7,6 +7,7 @@ import httpx import pytest +from ucode import custom_oauth from ucode.agents import copilot WS = "https://example.databricks.com" @@ -377,7 +378,7 @@ def test_needs_no_token_and_writes_none(self, tmp_path, monkeypatch): def no_token(*args, **kwargs): raise AssertionError("configure must not mint a token") - monkeypatch.setattr(copilot, "get_databricks_token", no_token) + monkeypatch.setattr(custom_oauth, "get_databricks_token", no_token) copilot.write_tool_config({"workspace": WS}, "gpt-5") @@ -480,7 +481,7 @@ def get_token(workspace, profile=None, *, force_refresh=False): calls.append((workspace, profile, force_refresh)) return "profile-token" - monkeypatch.setattr(copilot, "get_databricks_token", get_token) + monkeypatch.setattr(custom_oauth, "get_databricks_token", get_token) token = copilot._token_provider({"workspace": WS, "profile": "dev"})(True) @@ -494,7 +495,7 @@ def get_custom(workspace, client_id, redirect_url, *, scopes, profile, force_ref calls.append((workspace, client_id, redirect_url, scopes, profile, force_refresh)) return "custom-token" - monkeypatch.setattr(copilot, "get_custom_client_token", get_custom) + monkeypatch.setattr(custom_oauth, "get_custom_client_token", get_custom) state = { "workspace": WS, "custom_oauth": { @@ -524,7 +525,7 @@ def wait(self): def _stub_launch(monkeypatch, process=None, token="dbx-token"): """Run launch against the real proxy with a stubbed Copilot process; return what it saw.""" seen = {} - monkeypatch.setattr(copilot, "get_databricks_token", lambda *args, **kwargs: token) + monkeypatch.setattr(custom_oauth, "get_databricks_token", lambda *args, **kwargs: token) monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") def popen(argv, *, env): @@ -617,7 +618,7 @@ def test_ctrl_c_keeps_waiting_for_copilot(self, tmp_path, monkeypatch): def test_proxy_stops_when_copilot_cannot_start(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) - monkeypatch.setattr(copilot, "get_databricks_token", lambda *args, **kwargs: "tok") + monkeypatch.setattr(custom_oauth, "get_databricks_token", lambda *args, **kwargs: "tok") monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") seen = {} diff --git a/tests/test_cli.py b/tests/test_cli.py index 81acc425b..9462d2236 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3444,7 +3444,7 @@ def test_reuses_cli_proxy_and_auth_for_the_configured_workspace(self): patch( "ucode.gateway_proxy.start_otel_proxy", return_value=(server, cache, client) ) as start, - patch("ucode.agents.codex.get_databricks_token", return_value="fresh-token") as token, + patch("ucode.custom_oauth.get_databricks_token", return_value="fresh-token") as token, ): result = runner.invoke(app, ["otel-proxy", "--host", workspace, "--port", "43199"]) assert start.call_args.args[1](True) == "fresh-token" diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 4c26cdaad..b48fbdac3 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -6,7 +6,7 @@ import pytest -from ucode import codex_config +from ucode import codex_config, custom_oauth from ucode.agents import LaunchOptions, codex from ucode.smart_routing import codex_interposer, codex_routing, v2 @@ -97,7 +97,7 @@ def test_codex_launch_bypasses_routing_for_other_shapes(self, tmp_path, monkeypa monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) monkeypatch.setattr(codex, "agent_version", lambda binary: "0.144.0") - monkeypatch.setattr(codex, "get_databricks_token", lambda *_args, **_kw: "token") + monkeypatch.setattr(custom_oauth, "get_databricks_token", lambda *_args, **_kw: "token") monkeypatch.setattr(v2, "launch_codex", lambda *args, **kwargs: pytest.fail("launched")) monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index da89d37f6..92d86a216 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -1134,6 +1134,7 @@ def test_launch_copilot_per_model( finally: cache.stop() server.shutdown() + server.server_close() client.close() assert not failures, "Copilot launch failures:\n" + "\n".join(failures) From d33347593ce5c9fe0d62b2b1162567fd317ee705 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Tue, 20 Oct 2026 11:21:27 +0000 Subject: [PATCH 07/11] copilot: discover and switch between all authorized models Co-authored-by: Isaac --- README.md | 25 +- src/ucode/agents/__init__.py | 4 +- src/ucode/agents/args.py | 1 + src/ucode/agents/copilot.py | 621 ++++++++++++--- src/ucode/cli.py | 45 +- src/ucode/databricks.py | 72 +- src/ucode/doctor.py | 23 +- src/ucode/managed_resolve.py | 9 +- src/ucode/managed_setup.py | 7 +- tests/AGENTS.md | 2 +- tests/README.md | 5 +- tests/integration/README.md | 10 +- tests/integration/test_ug_copilot_headless.py | 28 +- tests/integration/test_ug_copilot_models.py | 165 ++++ tests/integration/utils/agents/copilot.py | 48 ++ tests/integration/utils/constants.py | 2 + tests/integration/utils/evidence.py | 19 +- tests/integration/utils/terminal.py | 13 +- tests/test_agent_copilot.py | 751 +++++++++++------- tests/test_agents_init.py | 9 +- tests/test_cli.py | 81 ++ tests/test_databricks.py | 28 + tests/test_doctor.py | 27 + tests/test_e2e.py | 34 +- tests/test_managed_resolve.py | 16 + tests/test_managed_setup.py | 9 +- 26 files changed, 1575 insertions(+), 479 deletions(-) create mode 100644 tests/integration/test_ug_copilot_models.py create mode 100644 tests/integration/utils/agents/copilot.py diff --git a/README.md b/README.md index fe7f7b0a0..deccce8ee 100644 --- a/README.md +++ b/README.md @@ -68,17 +68,24 @@ model for one launch. OpenCode's `provider/model` form is also accepted. Unknown Databricks models produce an error; this option does not add models to discovery or change ug's saved default. -`ug copilot` uses the Anthropic Messages API for Claude model IDs, the Responses -API for GPT-5 and newer model IDs, and Chat Completions for other models. The -model selected at launch (including `--model`) determines the API. An inherited `COPILOT_PROVIDER_WIRE_MODEL` takes precedence -because it overrides the model sent to the gateway. Restart Copilot through `ug` -to change the wire model or API; in-session model selection does not rebuild its -provider configuration. +`ug copilot` lists every Claude and GPT model you can use on the workspace in +Copilot's `/model` picker, and you can switch between them mid-session. Claude +models use the Anthropic Messages API, GPT-5 and newer use Responses, and older +GPT models use Chat Completions. Gemini and open-source models aren't offered. +`ug copilot --model claude-opus-5-5` (or `system.ai.claude-opus-5-5`, or Copilot's +own `dbx-anthropic/system.ai.claude-opus-5-5`) picks the model a session starts on; +an unknown or unsupported model fails before Copilot starts and lists the +available ones. ug passes Copilot a per-session model registry, so your own +`~/.copilot/providers.json` and `COPILOT_PROVIDER_*` settings apply only to plain +`copilot`. Use plain `/model`, not `/model --global`, inside `ug copilot`: a +global choice is saved for plain `copilot` too, where the Databricks models don't +exist. Requires Copilot CLI 1.0.91 or newer. `ug copilot` runs Copilot against a loopback proxy for the session. Copilot holds only -a random session secret; the proxy adds a fresh Databricks token to each gateway -request, so OAuth sessions keep working past token expiry without a restart, and no -token is written to `~/.copilot/ucode.env`. If the Databricks sign-in itself expires, +a random session secret, kept in the session registry under `~/.ucode/run/` (readable +only by you, deleted when the session ends); the proxy adds a fresh Databricks token to +each gateway request, so OAuth sessions keep working past token expiry without a +restart, and no Databricks token is written to disk. If the Databricks sign-in itself expires, Copilot shows a "Databricks sign-in expired" error; run `databricks auth login` and restart `ug copilot`. Proxy errors are logged to `~/.ucode/logs/copilot-proxy.log`. diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 3025ac606..34f7d307b 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -511,7 +511,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: if tool == "gemini": return bool(state.get("gemini_models")) if tool == "copilot": - return bool(state.get("claude_models")) or bool(state.get("codex_models")) + return bool(copilot.copilot_catalog(state)) if tool == "pi": return ( bool(state.get("claude_models")) @@ -527,7 +527,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: "kilo": ("claude", "gemini", "oss"), "codex": ("codex",), "gemini": ("gemini",), - "copilot": ("claude", "codex"), + "copilot": ("copilot",), "pi": ("claude", "codex", "gemini"), } diff --git a/src/ucode/agents/args.py b/src/ucode/agents/args.py index faca49dfe..2de251eb9 100644 --- a/src/ucode/agents/args.py +++ b/src/ucode/agents/args.py @@ -11,6 +11,7 @@ class LaunchOptions: launch_smart_routing: bool = False user_pinned_model: str | None = None + launch_model: str | None = None def explicit_model_arg_value(tool_args: list[str]) -> str | None: diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index d08f7eba5..d3f443da3 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -1,37 +1,37 @@ -"""GitHub Copilot CLI agent: writes ~/.copilot/ucode.env and injects BYOK env vars at launch. - -Copilot CLI's BYOK config (COPILOT_PROVIDER_*) is documented as env-var only — -the CLI does not auto-load ~/.copilot/ucode.env. We still write the file so users can -inspect what's configured and to give `revert` something to clean up; it never holds -a credential. - -Copilot holds its provider credential for the whole session, so `ug copilot` runs it -against a per-session loopback proxy (`gateway_proxy.start_gateway_auth_proxy`): -Copilot authenticates to the proxy with a random session secret, and the proxy -forwards each request with a fresh Databricks token. - -System Claude models (`system.ai.claude-*`) use Copilot CLI's `anthropic` provider against the gateway's -Anthropic endpoint, with `COPILOT_PROVIDER_MODEL_ID` set to Copilot's catalog id -for the model so Copilot sends Claude request settings. Other models use the -`openai` provider against the Databricks MLflow gateway: GPT models with major -version 5 or newer use Responses, because Copilot sends `temperature: 0` on Chat -Completions and those models accept only the default; the rest use Chat Completions. Copilot fixes its wire API and model when it builds the native -client, so changing models in the picker cannot change either mid-session. -Relaunch Copilot after changing model families. Gemini is intentionally excluded -— Databricks' Gemini translation layer rejects the `stream_options` field that -Copilot CLI sends, so Gemini models 400 on every request. +"""GitHub Copilot CLI agent: runs Copilot against a per-session Databricks model registry. + +Copilot reads its BYOK models from a providers.json registry (`COPILOT_PROVIDERS_CONFIG`). Each +`ug copilot` launch writes one under ~/.ucode/run/ listing every Databricks model Copilot can use, +so `/model` switches between them mid-session, across wire APIs. Each family keeps its native +wire: Claude on the gateway's Anthropic endpoint (with Copilot's catalog id as `modelId`, so +Copilot sends Claude request settings), GPT 5+ on Responses (Copilot sends `temperature: 0` on Chat +Completions, which those models reject), older GPT on Chat Completions. Gemini and open-source +models are excluded: Databricks' Gemini translation rejects the `stream_options` field Copilot +sends, and open-source models are not validated with Copilot yet. + +Copilot holds its provider credential for the whole session, so every provider points at a +per-session loopback proxy (`gateway_proxy.start_gateway_auth_proxy`): Copilot authenticates to +the proxy with a random session secret, written only into the session registry, and the proxy +forwards each request with a fresh Databricks token. ~/.copilot/ucode.env is written for +inspection and `revert`; Copilot never reads it and it never holds a credential. """ from __future__ import annotations +import json import os import re import secrets +import shutil +import tempfile import threading import time -from collections.abc import Callable +from collections.abc import Callable, Iterable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import UTC, datetime from pathlib import Path -from typing import Protocol +from typing import IO, Any, Protocol from ucode.config_io import ( APP_DIR, @@ -46,14 +46,22 @@ from ucode.constants import AI_GATEWAY_REQUEST_TAGS_HEADER, LOOPBACK_HOST from ucode.custom_oauth import launch_token from ucode.databricks import ( + ANTHROPIC_FAMILIES, + ModelServiceEntry, build_copilot_base_url, build_tool_base_url, + model_version_sort_key, ) from ucode.os_compatibility import subprocess_cross_os +from ucode.os_compatibility.file_lock_cross_os import ( + acquire_exclusive_file_lock, + release_file_lock, +) from ucode.state import is_tool_managed, mark_tool_managed, save_state from ucode.telemetry import agent_version, ug_version +from ucode.ui import print_warning -from .args import LaunchOptions, explicit_model_arg_value +from .args import LaunchOptions from .request_tags import request_tags_header_value COPILOT_CONFIG_DIR = Path.home() / ".copilot" @@ -62,6 +70,13 @@ COPILOT_BACKUP_PATH = APP_DIR / "copilot-ucode-env.backup" COPILOT_MCP_BACKUP_PATH = APP_DIR / "copilot-ucode-mcp-config.backup.json" COPILOT_PROXY_LOG_PATH = APP_DIR / "logs" / "copilot-proxy.log" +COPILOT_RUN_DIR = APP_DIR / "run" +_SESSION_DIR_PREFIX = "copilot-" +_SESSION_LOCK_NAME = "session.lock" +_RUN_DIR_LOCK_NAME = ".copilot.lock" +_REGISTRY_NAME = "providers.json" + +MINIMUM_COPILOT_VERSION = (1, 0, 91) SPEC: ToolSpec = { "binary": "copilot", @@ -71,45 +86,110 @@ "backup_path": COPILOT_BACKUP_PATH, } -MANAGED_KEYS: list[str] = [ - "COPILOT_PROVIDER_TYPE", - "COPILOT_PROVIDER_BASE_URL", - "COPILOT_PROVIDER_WIRE_API", - "COPILOT_PROVIDER_MODEL_ID", - "COPILOT_MODEL", - "COPILOT_OFFLINE", -] +MANAGED_KEYS: list[str] = ["COPILOT_MODEL", "COPILOT_OFFLINE"] LEGACY_ENV_KEYS = [ "OPENAI_BASE_URL", "OPENAI_API_KEY", "COPILOT_PROVIDER_API_KEY", + # Single-provider settings earlier versions wrote; the session registry replaces them. + "COPILOT_PROVIDER_TYPE", + "COPILOT_PROVIDER_BASE_URL", + "COPILOT_PROVIDER_WIRE_API", + "COPILOT_PROVIDER_MODEL_ID", ] # Earlier versions wrote the Databricks token (or PAT) under these keys. CREDENTIAL_ENV_KEYS = ("COPILOT_PROVIDER_BEARER_TOKEN", "OAUTH_TOKEN") -# Inherited settings that would bypass the proxy's credential or its HTTP transport. -_INHERITED_PROVIDER_OVERRIDES = ( - "COPILOT_PROVIDER_API_KEY", - "COPILOT_PROVIDER_API_KEY_COMMAND", - "COPILOT_PROVIDER_TRANSPORT", -) -_GPT_MODEL_MAJOR_PATTERN = re.compile(r"^(?:system\.ai\.)?(?:databricks-)?gpt-(\d+)(?=$|[.-])") +CATALOG_STATE_KEY = "copilot_catalog" - -_SYSTEM_CLAUDE_MODEL_PATTERN = re.compile(r"^system\.ai\.(claude-[^\[]+)(?:\[[^\]]*\])?$") +PROVIDER_NAMES: dict[str, str] = { + "anthropic": "dbx-anthropic", + "responses": "dbx-responses", + "completions": "dbx-completions", +} +_WIRE_ROUTES: dict[str, tuple[str, ...]] = { + "anthropic": ("anthropic/v1/messages",), + "responses": ("mlflow/v1/responses", "openai/v1/responses", "codex/v1/responses"), + "completions": ("mlflow/v1/chat/completions", "openai/v1/chat/completions"), +} +_UNSUPPORTED_MODELS = { + "claude-opus-4-1": "the gateway rejects the thinking settings Copilot sends for it", +} +_GATEWAY_ALIAS_PREFIX = "anthropic-aigw-" +_GPT_MODEL_MAJOR_PATTERN = re.compile(r"^gpt-(\d+)(?=$|[.-])") _CLAUDE_DASHED_VERSION_PATTERN = re.compile(r"-(\d+)-(\d+)$") _CLAUDE_HAIKU_MAJOR_PATTERN = re.compile(r"^claude-haiku-(\d+)") +_CLAUDE_CATALOG_VERSION_PATTERN = re.compile(r"^claude-([a-z]+)-(\d+)(?:\.(\d+))?$") + +# (maxPromptTokens, maxOutputTokens) by Copilot catalog id. Copilot assumes a 128K prompt for +# registry models without limits, and its own values differ between versions. +_CLAUDE_LIMITS: dict[str, tuple[int, int]] = { + "claude-fable-5.1": (936_000, 32_000), + "claude-fable-5": (936_000, 32_000), + "claude-opus-5.5": (936_000, 32_000), + "claude-opus-5": (936_000, 32_000), + "claude-opus-4.8": (936_000, 32_000), + "claude-opus-4.7": (936_000, 32_000), + "claude-opus-4.6": (936_000, 32_000), + "claude-opus-4.5": (168_000, 32_000), + "claude-sonnet-5.5": (872_000, 128_000), + "claude-sonnet-5": (936_000, 32_000), + "claude-sonnet-4.6": (936_000, 32_000), + "claude-sonnet-4.5": (168_000, 32_000), + "claude-haiku-4.5": (136_000, 64_000), +} +_CLAUDE_1M_LIMITS = (936_000, 32_000) +_CLAUDE_200K_LIMITS = (168_000, 32_000) + + +def _model_tail(model: str) -> str: + """The model name without its ``..`` and ``databricks-`` prefixes or ``[...]`` suffix.""" + tail = model.split("[", 1)[0].rsplit(".", 1)[-1].lower() + return tail.removeprefix("databricks-") + + +def copilot_exclusion_reason(model: str, api_types: Iterable[str] = ()) -> str | None: + """Why ``model`` can't be offered in Copilot, or None when it can.""" + return _classify(model, tuple(api_types))[1] + + +def classify_copilot_model(model: str, api_types: Iterable[str] = ()) -> str | None: + """The wire API (``anthropic``, ``responses`` or ``completions``) Copilot uses for ``model``. + + Names decide the wire; ``api_types`` (the model service's advertised routes) only excludes, + because the gateway serves some models on routes they don't list.""" + return _classify(model, tuple(api_types))[0] + + +def _classify(model: str, api_types: tuple[str, ...]) -> tuple[str | None, str | None]: + if model.lower().startswith(_GATEWAY_ALIAS_PREFIX): + return None, "it is a gateway alias, not a model" + tail = _model_tail(model) + if tail in _UNSUPPORTED_MODELS: + return None, _UNSUPPORTED_MODELS[tail] + gpt = _GPT_MODEL_MAJOR_PATTERN.match(tail) + if tail.startswith("claude-"): + wire = "anthropic" + elif gpt: + wire = "responses" if int(gpt.group(1)) >= 5 else "completions" + elif "gemini-" in tail: + return None, "the gateway rejects the requests Copilot sends to Gemini models" + else: + return None, "ug copilot supports Claude and GPT models" + if api_types and not set(api_types) & set(_WIRE_ROUTES[wire]): + return None, f"it isn't served on {_WIRE_ROUTES[wire][0]}" + return wire, None def copilot_catalog_model_id(model: str) -> str | None: - """Copilot's catalog id for a system Claude model, e.g. ``system.ai.claude-sonnet-5-5`` -> ``claude-sonnet-5.5``. + """Copilot's catalog id for a Claude model, e.g. ``system.ai.claude-sonnet-5-5`` -> ``claude-sonnet-5.5``. Copilot applies Claude request settings to any ``claude-*`` id, so ids missing from its catalog still work. - Returns None for anything other than ``system.ai.claude-*``. + Returns None for anything other than Claude. """ - match = _SYSTEM_CLAUDE_MODEL_PATTERN.match(model) - if match is None: + tail = _model_tail(model) + if not tail.startswith("claude-"): return None - catalog_id = _CLAUDE_DASHED_VERSION_PATTERN.sub(r"-\1.\2", match.group(1)) + catalog_id = _CLAUDE_DASHED_VERSION_PATTERN.sub(r"-\1.\2", tail) # Copilot sends budget thinking for any Haiku it doesn't know, which Haiku 5+ rejects; Sonnet # ids get the adaptive thinking those models require. haiku = _CLAUDE_HAIKU_MAJOR_PATTERN.match(catalog_id) @@ -118,74 +198,368 @@ def copilot_catalog_model_id(model: str) -> str | None: return catalog_id -def model_uses_responses_api(model: str) -> bool: - """Whether a supported GPT model id uses the Responses API.""" - match = _GPT_MODEL_MAJOR_PATTERN.match(model) - return match is not None and int(match.group(1)) >= 5 +def copilot_model_limits(model: str) -> tuple[int, int] | None: + """``(maxPromptTokens, maxOutputTokens)`` for a Claude model, or None for other models.""" + catalog_id = copilot_catalog_model_id(model) + if catalog_id is None: + return None + if catalog_id in _CLAUDE_LIMITS: + return _CLAUDE_LIMITS[catalog_id] + version = _CLAUDE_CATALOG_VERSION_PATTERN.match(catalog_id) + if version is None: + return _CLAUDE_200K_LIMITS + family, major, minor = version.group(1), int(version.group(2)), int(version.group(3) or 0) + one_million = major >= 5 or (family in ("opus", "sonnet") and (major, minor) >= (4, 6)) + return _CLAUDE_1M_LIMITS if one_million else _CLAUDE_200K_LIMITS + + +@dataclass(frozen=True) +class CopilotModel: + """A Databricks model in Copilot's registry and the wire API it is served on.""" + + id: str + wire: str + + @property + def provider(self) -> str: + return PROVIDER_NAMES[self.wire] + + @property + def ref(self) -> str: + """Copilot's selection id; a bare model id isn't matched against the registry.""" + return f"{self.provider}/{self.id}" + + +def _catalog_sort_key(model: CopilotModel) -> tuple: + family = next( + ( + i + for i, f in enumerate(ANTHROPIC_FAMILIES) + if _model_tail(model.id).startswith(f"claude-{f}-") + ), + len(ANTHROPIC_FAMILIES), + ) + return (tuple(PROVIDER_NAMES).index(model.wire), family, model_version_sort_key(model.id)) + + +def build_copilot_catalog(entries: Iterable[ModelServiceEntry | str]) -> list[CopilotModel]: + """The models Copilot can use, Claude first (by family, newest first), then GPT.""" + models: dict[str, CopilotModel] = {} + for entry in entries: + model, api_types = (entry, ()) if isinstance(entry, str) else (entry.id, entry.api_types) + wire = classify_copilot_model(model, api_types) + if wire is not None and model not in models: + models[model] = CopilotModel(model, wire) + return sorted(models.values(), key=_catalog_sort_key) -def default_model(state: dict) -> str | None: - """Prefer Claude sonnet, then opus/haiku, then codex. +def catalog_from_listings( + entries: Iterable[ModelServiceEntry] | None, fallback: Iterable[str] +) -> list[CopilotModel] | None: + """The Copilot catalog from this run's model listings. - A managed config's ``copilot_default_model`` and ``copilot_models`` both win outright: the former is - the admin's chosen session start, the latter their allowlist. Workspace-wide discovery falls back. - """ - if isinstance(state.get("copilot_default_model"), str): - return state.get("copilot_default_model") - copilot_models = state.get("copilot_models") or [] - if isinstance(copilot_models, list) and copilot_models: - return copilot_models[0] - claude_models = state.get("claude_models") or {} - for family in ("sonnet", "opus", "haiku"): - if claude_models.get(family): - return claude_models[family] - codex_models = state.get("codex_models") or [] - if codex_models: - return codex_models[0] - return next(iter(claude_models.values()), None) - - -def render_env_overlay( - base_url: str, - selected_model: str, - *, - override_model: str | None = None, -) -> dict[str, str]: - """Copilot's provider settings for ``selected_model``, rooted at ``base_url``. - - ``base_url`` is the workspace for the inspectable env file and the loopback proxy - at launch. Credentials are added only to the launch environment.""" - request_model = override_model or selected_model - catalog_model_id = copilot_catalog_model_id(request_model) - if catalog_model_id is not None: - return { - "COPILOT_PROVIDER_TYPE": "anthropic", - "COPILOT_PROVIDER_BASE_URL": build_tool_base_url("claude", base_url), - "COPILOT_PROVIDER_MODEL_ID": catalog_model_id, - "COPILOT_MODEL": selected_model, - "COPILOT_OFFLINE": "true", - } - wire_api = "responses" if model_uses_responses_api(request_model) else "completions" + ``entries`` is the UC model-services listing (None when it failed); ``fallback`` holds the ids + the per-family gateway listings found instead. None when neither answered, so the caller keeps + the saved catalog; an empty catalog when they answered with nothing Copilot can use.""" + listed = list(entries or ()) + if listed: + return build_copilot_catalog(listed) + fallback_ids = list(fallback) + return build_copilot_catalog(fallback_ids) if fallback_ids else None + + +def catalog_state(models: list[CopilotModel]) -> dict: return { - "COPILOT_PROVIDER_TYPE": "openai", - "COPILOT_PROVIDER_BASE_URL": build_copilot_base_url(base_url), - "COPILOT_PROVIDER_WIRE_API": wire_api, - "COPILOT_MODEL": selected_model, - "COPILOT_OFFLINE": "true", + "models": [model.id for model in models], + "discovered_at": datetime.now(UTC).isoformat(timespec="seconds"), } -def build_runtime_env(proxy_base_url: str, model: str, session_secret: str) -> dict[str, str]: - """Copilot's launch environment: provider settings pointed at the session proxy.""" - env = os.environ.copy() - for key in _INHERITED_PROVIDER_OVERRIDES: - env.pop(key, None) - override_model = env.get("COPILOT_PROVIDER_WIRE_MODEL") - env.update(render_env_overlay(proxy_base_url, model, override_model=override_model)) - env["COPILOT_PROVIDER_BEARER_TOKEN"] = session_secret +def _string_list(value: object) -> list[str]: + return ( + [item for item in value if isinstance(item, str) and item] + if isinstance(value, list) + else [] + ) + + +def copilot_catalog(state: dict) -> list[CopilotModel]: + """The models this launch registers: the managed list, else the discovered catalog. + + State saved before the catalog existed falls back to the Claude and GPT discovery lists.""" + static = build_copilot_catalog(_string_list(state.get("copilot_static_models"))) + if static: + return static + saved = state.get(CATALOG_STATE_KEY) + if isinstance(saved, dict): + return build_copilot_catalog(_string_list(saved.get("models"))) + claude_models = state.get("claude_models") + legacy = list(claude_models.values()) if isinstance(claude_models, dict) else [] + return build_copilot_catalog([*_string_list(legacy), *_string_list(state.get("codex_models"))]) + + +def unservable_static_models(state: dict) -> list[tuple[str, str]]: + """Managed models Copilot can't use, with the reason, so the launch can say why they're missing. + + Empty when none of the list is usable: the launch then reports that the whole list was + ignored instead.""" + static = _string_list(state.get("copilot_static_models")) + if not build_copilot_catalog(static): + return [] + skipped = [] + for model in static: + reason = copilot_exclusion_reason(model) + if reason is not None: + skipped.append((model, reason)) + return skipped + + +def _normalized(model: str) -> str: + return model.split("[", 1)[0].lower() + + +def resolve_copilot_model(requested: str, catalog: list[CopilotModel]) -> CopilotModel: + """Match ``requested`` (an id, a short name or ``provider/id``) to a catalog model. + + Raises RuntimeError listing the available models when nothing matches.""" + for matches in ( + lambda m: m.id == requested, + lambda m: m.ref == requested, + lambda m: _normalized(m.id) == _normalized(requested), + lambda m: _normalized(m.id) == f"system.ai.{_normalized(requested)}", + lambda m: _normalized(m.id) == f"databricks-{_normalized(requested)}", + ): + found = [model for model in catalog if matches(model)] + if found: + return next((m for m in found if m.id.startswith("system.ai.")), found[0]) + reason = copilot_exclusion_reason(requested) + detail = f" ({reason})" if reason else "" + available = "\n".join( + f" {provider}: {', '.join(m.id for m in catalog if m.provider == provider)}" + for provider in PROVIDER_NAMES.values() + if any(m.provider == provider for m in catalog) + ) + raise RuntimeError( + f"Model '{requested}' isn't available in ug copilot{detail}. Available models:\n{available}" + ) + + +def _first_available( + candidates: Iterable[str | None], catalog: list[CopilotModel], *, warn: bool +) -> CopilotModel | None: + for candidate in candidates: + if not candidate: + continue + try: + return resolve_copilot_model(candidate, catalog) + except RuntimeError: + if warn: + print_warning( + f"{candidate} isn't available in ug copilot; starting on another model." + ) + return None + + +def launch_model( + state: dict, catalog: list[CopilotModel], *, requested: str | None = None, warn: bool = True +) -> CopilotModel: + """The model a session starts on: ``requested`` (budget recommendation or managed default), + then the managed default and list, then the newest Sonnet, Opus or Haiku, then GPT.""" + static = _string_list(state.get("copilot_static_models")) + chosen = _first_available( + [requested, state.get("copilot_default_model"), static[0] if static else None], + catalog, + warn=warn, + ) + if chosen is not None: + return chosen + for prefix in ("claude-sonnet-", "claude-opus-", "claude-haiku-"): + family = next((m for m in catalog if _model_tail(m.id).startswith(prefix)), None) + if family is not None: + return family + return next((m for m in catalog if m.wire != "anthropic"), catalog[0]) + + +def default_model(state: dict) -> str | None: + catalog = copilot_catalog(state) + return launch_model(state, catalog, warn=False).id if catalog else None + + +def rewrite_model_args( + tool_args: list[str], catalog: list[CopilotModel] +) -> tuple[list[str], CopilotModel | None]: + """Resolve every ``--model``/``-m`` before ``--`` to Copilot's ``provider/id`` form. + + Copilot has no ``-m`` alias, so it becomes ``--model``. Returns the rewritten args and the + last model selected, which is the one Copilot uses.""" + rewritten: list[str] = [] + selected: CopilotModel | None = None + index = 0 + while index < len(tool_args): + arg = tool_args[index] + if arg == "--": + rewritten.extend(tool_args[index:]) + break + value: str | None = None + if ( + arg in {"--model", "-m"} + and index + 1 < len(tool_args) + and not tool_args[index + 1].startswith("-") + ): + value = tool_args[index + 1] + index += 1 + elif arg.startswith("--model=") and arg.partition("=")[2]: + value = arg.partition("=")[2] + if value is None: + rewritten.append(arg) + else: + selected = resolve_copilot_model(value, catalog) + rewritten.extend(["--model", selected.ref]) + index += 1 + return rewritten, selected + + +def render_providers_config( + catalog: list[CopilotModel], proxy_base_url: str, session_secret: str +) -> dict: + """Copilot's providers.json for ``catalog``: one provider per wire, all on the session proxy. + + A provider with no models makes Copilot probe ``GET /models``, so only used wires are + listed.""" + providers = [] + for wire, name in PROVIDER_NAMES.items(): + if not any(model.wire == wire for model in catalog): + continue + if wire == "anthropic": + provider = { + "name": name, + "type": "anthropic", + "baseUrl": build_tool_base_url("claude", proxy_base_url), + } + else: + provider = { + "name": name, + "type": "openai", + "wireApi": wire, + "baseUrl": build_copilot_base_url(proxy_base_url), + } + # `bearerToken` reaches the proxy as `Authorization` on every provider type. + provider["bearerToken"] = session_secret + providers.append(provider) + models = [] + for model in catalog: + entry: dict[str, Any] = { + "id": model.id, + "provider": model.provider, + "wireModel": model.id, + "name": model.id, + } + catalog_id = copilot_catalog_model_id(model.id) if model.wire == "anthropic" else None + if catalog_id is not None: + entry["modelId"] = catalog_id + limits = copilot_model_limits(model.id) if model.wire == "anthropic" else None + if limits is not None: + entry["maxPromptTokens"], entry["maxOutputTokens"] = limits + models.append(entry) + return {"providers": providers, "models": models} + + +def build_runtime_env(registry_path: Path, model: CopilotModel) -> dict[str, str]: + """Copilot's launch environment: the session registry and the model to start on. + + Inherited ``COPILOT_PROVIDER*`` settings are dropped so a user's own BYOK setup can't take + over a ug session.""" + env = { + key: value for key, value in os.environ.items() if not key.startswith("COPILOT_PROVIDER") + } + env["COPILOT_PROVIDERS_CONFIG"] = str(registry_path) + env["COPILOT_MODEL"] = model.ref + # Fail closed without the registry instead of falling back to GitHub sign-in. + env["COPILOT_OFFLINE"] = "true" return env +@contextmanager +def _run_dir_lock() -> Iterator[None]: + COPILOT_RUN_DIR.mkdir(parents=True, exist_ok=True) + with (COPILOT_RUN_DIR / _RUN_DIR_LOCK_NAME).open("a+b") as lock_file: + acquire_exclusive_file_lock(lock_file) + try: + yield + finally: + release_file_lock(lock_file) + + +def _sweep_dead_sessions() -> None: + for session_dir in COPILOT_RUN_DIR.glob(f"{_SESSION_DIR_PREFIX}*"): + lock_path = session_dir / _SESSION_LOCK_NAME + if lock_path.exists(): + with lock_path.open("a+b") as lock_file: + try: + acquire_exclusive_file_lock(lock_file, timeout=0) + except (TimeoutError, OSError): + continue + release_file_lock(lock_file) + shutil.rmtree(session_dir, ignore_errors=True) + + +def sweep_session_dirs() -> None: + """Delete registries left by sessions that ended without cleaning up (e.g. a killed terminal).""" + if not COPILOT_RUN_DIR.exists(): + return + with _run_dir_lock(): + _sweep_dead_sessions() + + +@contextmanager +def session_registry(config: dict) -> Iterator[Path]: + """Write ``config`` to a private per-session providers.json, removed when the session ends. + + The session holds a lock on its directory so a concurrent sweep leaves it alone.""" + lock_file: IO[Any] | None = None + with _run_dir_lock(): + _sweep_dead_sessions() + session_dir = Path(tempfile.mkdtemp(prefix=_SESSION_DIR_PREFIX, dir=COPILOT_RUN_DIR)) + try: + lock_file = (session_dir / _SESSION_LOCK_NAME).open("a+b") + acquire_exclusive_file_lock(lock_file, timeout=0) + path = session_dir / _REGISTRY_NAME + descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(descriptor, "w", encoding="utf-8") as registry: + json.dump(config, registry, indent=2) + except BaseException: + if lock_file is not None: + lock_file.close() + shutil.rmtree(session_dir, ignore_errors=True) + raise + try: + yield path + finally: + release_file_lock(lock_file) + lock_file.close() + shutil.rmtree(session_dir, ignore_errors=True) + + +def _parse_version(value: str) -> tuple[int, int, int] | None: + match = re.search(r"(\d+)\.(\d+)\.(\d+)", value) + if match is None: + return None + major, minor, patch = match.groups() + return int(major), int(minor), int(patch) + + +def minimum_version_error() -> str | None: + installed = agent_version(SPEC["binary"]) + parsed = _parse_version(installed) + if parsed is None or parsed >= MINIMUM_COPILOT_VERSION: + return None + required = ".".join(str(part) for part in MINIMUM_COPILOT_VERSION) + return ( + f"GitHub Copilot CLI {installed} is too old. ucode requires {required} or newer for " + f"Databricks model switching. Update it with `npm install -g {SPEC['package']}`." + ) + + def build_mcp_server_entry(argv: list[str]) -> dict: # A `local` MCP server runs a stdio command; `command`/`args` split the # argv. ug registers the `ug mcp-proxy ...` bridge here so Copilot @@ -250,16 +624,13 @@ def write_tool_config(state: dict, model: str) -> dict: if not is_tool_managed(state, "copilot"): backup_existing_file(COPILOT_ENV_PATH, COPILOT_BACKUP_PATH) existing = parse_dotenv(COPILOT_ENV_PATH) - # Keep the inspectable file self-consistent without treating it as launch input. - override_model = existing.get("COPILOT_PROVIDER_WIRE_MODEL") - overlay = render_env_overlay(state["workspace"], model, override_model=override_model) for key in LEGACY_ENV_KEYS: existing.pop(key, None) - # Wire API is openai-only and the catalog id is anthropic-only; drop whichever this model doesn't use. - for key in ("COPILOT_PROVIDER_WIRE_API", "COPILOT_PROVIDER_MODEL_ID"): - if key not in overlay: - existing.pop(key, None) - existing.update(overlay) + try: + existing["COPILOT_MODEL"] = resolve_copilot_model(model, copilot_catalog(state)).ref + except RuntimeError: + existing["COPILOT_MODEL"] = model + existing["COPILOT_OFFLINE"] = "true" write_dotenv(COPILOT_ENV_PATH, existing) state = mark_tool_managed(state, "copilot", MANAGED_KEYS) save_state(state) @@ -313,9 +684,17 @@ def _wait_for_exit(proc: _ChildProcess) -> int: def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: - model = explicit_model_arg_value(tool_args) or options.user_pinned_model or default_model(state) - if not model: - raise RuntimeError("No Copilot model is available on this workspace.") + catalog = copilot_catalog(state) + if not catalog: + raise RuntimeError( + "No Copilot model is available on this workspace: ug copilot supports its Claude and " + "GPT models. Run `ug configure --agents copilot` to rediscover them." + ) + for model, reason in unservable_static_models(state): + print_warning(f"Skipping {model} from your workspace's managed config: {reason}.") + # Resolve before the proxy starts so a bad model fails without starting Copilot. + tool_args, selected = rewrite_model_args(tool_args, catalog) + start_model = selected or launch_model(state, catalog, requested=options.launch_model) scrub_stored_credentials() # Only Copilot launches need the proxy (and its `httpx` import). @@ -334,9 +713,13 @@ def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None threading.Thread(target=server.serve_forever, daemon=True).start() try: proxy_base_url = f"http://{LOOPBACK_HOST}:{server.server_address[1]}" - env = build_runtime_env(proxy_base_url, model, session_secret) - proc = subprocess_cross_os.popen([SPEC["binary"], *mcp_config_args(), *tool_args], env=env) - returncode = _wait_for_exit(proc) + config = render_providers_config(catalog, proxy_base_url, session_secret) + 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 + ) + returncode = _wait_for_exit(proc) finally: cache.stop() server.shutdown() diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 7367890a9..a3882f542 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -57,6 +57,7 @@ SKILLS_MCP_MIN_DATABRICKS_CLI_VERSION, apply_pat_environment, build_shared_base_urls, + cached_model_service_entries, databricks_cli_installed, discover_claude_models, discover_codex_models, @@ -644,11 +645,14 @@ def configure_shared_state( # Codex smart routing can select OSS models such as GLM, so a Codex-only # configure must persist that discovered family too. want_oss = fetch_all or "opencode" in tools or "kilo" in tools or "codex" in tools + want_copilot = fetch_all or "copilot" in tools claude_reason: str | None = None gemini_reason: str | None = None codex_reason: str | None = None oss_reason: str | None = None + copilot_reason: str | None = None + copilot_models: list[copilot_agent.CopilotModel] | None = None claude_models = {} gemini_models = [] codex_models = [] @@ -689,6 +693,12 @@ def configure_shared_state( codex_models, codex_reason = discover_codex_models(workspace, token) if want_oss: oss_models, oss_reason = ms_oss, ms_reason + if want_copilot: + copilot_models = copilot_agent.catalog_from_listings( + cached_model_service_entries(workspace), + [*claude_models.values(), *codex_models], + ) + copilot_reason = ms_reason or claude_reason or codex_reason if claude_models: opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: @@ -711,6 +721,13 @@ def configure_shared_state( state["codex_models"] = codex_models if want_oss: state["oss_models"] = oss_models + if want_copilot and copilot_models is not None: + state[copilot_agent.CATALOG_STATE_KEY] = copilot_agent.catalog_state(copilot_models) + elif want_copilot and state.get(copilot_agent.CATALOG_STATE_KEY): + print_warning( + f"Copilot model discovery failed ({copilot_reason}); " + "using the models found last time." + ) if fetch_all or "opencode" in tools or "kilo" in tools: # Kilo is an OpenCode fork and consumes the same discovered view. state["opencode_models"] = opencode_models @@ -726,6 +743,7 @@ def configure_shared_state( "gemini": gemini_reason, "codex": codex_reason, "oss": oss_reason, + "copilot": copilot_reason, } return state @@ -1178,17 +1196,15 @@ def _model_values(value: object) -> list[str]: def _status_models(tool: str, state: dict) -> list[str]: """Return the effective model allow-list for one configured agent.""" static_models = _model_values(state.get(f"{tool}_static_models")) - if static_models: + if tool == "copilot": + models = [model.id for model in copilot_agent.copilot_catalog(state)] + elif static_models: models = static_models elif tool in ("claude", "codex", "gemini", "opencode"): models = _model_values(state.get(f"{tool}_models")) elif tool == "kilo": # Kilo shares OpenCode's discovery; its models live under opencode_models. models = _model_values(state.get("opencode_models")) - elif tool == "copilot": - models = _model_values(state.get("copilot_models")) or ( - _model_values(state.get("claude_models")) + _model_values(state.get("codex_models")) - ) elif tool == "pi": models = _model_values(state.get("pi_models")) or ( _model_values(state.get("claude_models")) @@ -1201,6 +1217,8 @@ def _status_models(tool: str, state: dict) -> list[str]: def _status_default_model(tool: str, state: dict, models: list[str]) -> str | None: + if tool == "copilot": + return copilot_agent.default_model(state) explicit = state.get(f"{tool}_default_model") if isinstance(explicit, str) and explicit: return explicit @@ -1251,6 +1269,13 @@ def _live_status_model_state(state: dict, tools: set[str]) -> tuple[dict, str]: if oss_models: opencode_models["oss"] = oss_models live["opencode_models"] = opencode_models + if "copilot" in tools: + copilot_models = copilot_agent.catalog_from_listings( + cached_model_service_entries(workspace), + [*claude_models.values(), *codex_models], + ) + if copilot_models is not None: + live[copilot_agent.CATALOG_STATE_KEY] = copilot_agent.catalog_state(copilot_models) live["_status_model_reasons"] = { family: reason or shared_reason for family, reason in reasons.items() } @@ -1396,6 +1421,7 @@ def revert() -> int: mcp_results = revert_mcp_configs(state) # Before restoring backups: one written by an earlier ug may hold a token. copilot_agent.scrub_stored_credentials() + copilot_agent.sweep_session_dirs() claude_managed_result = claude_agent.revert_managed_settings() codex_managed_result = codex_agent.revert_managed_config() @@ -2984,6 +3010,8 @@ def _child_owns_stdout(tool: str, tool_args: list[str]) -> bool: tool_args = tool_args[: tool_args.index("--")] if tool == "claude": return any(arg in {"-p", "--print"} for arg in tool_args) + if tool == "copilot": + return any(arg in {"-p", "--prompt"} or arg.startswith("--prompt=") for arg in tool_args) return tool == "codex" and any(arg in {"exec", "e", "app-server"} for arg in tool_args) @@ -3014,8 +3042,10 @@ def _launch_options( explicit_prompt: bool, user_pinned_model: str | None, provider: str | None, + launch_model: str | None = None, ) -> LaunchOptions: return LaunchOptions( + launch_model=launch_model, # Pinned models for providers are resolved above through the provider-specific launch path. user_pinned_model=user_pinned_model if provider is None else None, launch_smart_routing=( @@ -3058,6 +3088,9 @@ def _skip_model_discovery( (``family_defaults``). In the last case it is also harmful: discovered family defaults would overwrite that writer's on every launch, and each rewrite of the managed file costs a sudo prompt. """ + if tool == "copilot": + # Copilot registers its discovered catalog, so a managed default alone still needs it. + return False return ( bool(provider) or bool(managed_parent_schema) @@ -3538,6 +3571,8 @@ def _launch_tool( # initial/fallback model and still participates in a routed session. user_pinned_model=model or forwarded_model, provider=provider, + # Copilot picks its start model from its own catalog; the admin's or budget pick leads. + launch_model=managed_model if tool == "copilot" else None, ) print_success(f"Starting {TOOL_SPECS[tool]['display']}") with _smart_routing_v2_flag( diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 6286a0276..8c5321764 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2097,6 +2097,24 @@ def model_token_limits(model_id: str) -> dict[str, int] | None: return None +@dataclass(frozen=True) +class ModelServiceEntry: + """One ``system.ai.*`` model service and the API types it advertises. + + ``api_types`` is advisory: the gateway serves some models on routes they don't list, so callers + use it to exclude non-chat models, never to pick a route.""" + + id: str + api_types: tuple[str, ...] = () + + +def _model_service_api_types(service: dict) -> tuple[str, ...]: + api_types = service.get("supported_api_types") + if not isinstance(api_types, list): + return () + return tuple(t for t in api_types if isinstance(t, str)) + + def _model_service_id(service: dict) -> str | None: """Extract the `system.ai.` id from one model-service entry. @@ -2150,7 +2168,7 @@ def _get_model_services_page( # Claude ids), so a single `ucode setup` run would otherwise page it twice. Cached per process, not # persisted: a long-lived process is not a thing here, and a new model appearing mid-command is not # worth a second walk. Failures are never cached, so a transient error still retries. -_MODEL_SERVICES_CACHE: dict[str, list[str]] = {} +_MODEL_SERVICES_CACHE: dict[str, list[ModelServiceEntry]] = {} # Same idea for the Model Provider Service listing (a different endpoint). It is workspace-wide and # filtered per agent afterwards, so `ucode setup` would otherwise re-list it once per MPS-capable @@ -2159,6 +2177,12 @@ def _get_model_services_page( _MODEL_PROVIDER_SERVICES_CACHE: dict[tuple[str, str], list[dict]] = {} +def cached_model_service_entries(workspace: str) -> list[ModelServiceEntry] | None: + """The model services this process already listed for ``workspace``, without a request.""" + cached = _MODEL_SERVICES_CACHE.get(workspace) + return list(cached) if cached is not None else None + + def clear_model_services_cache() -> None: """Forget cached model-service listings (used by tests, and after a workspace switch).""" _MODEL_SERVICES_CACHE.clear() @@ -2183,16 +2207,32 @@ def list_model_services( max_pages: int = 100, use_cache: bool = True, ) -> tuple[list[str], str | None]: - """List all `system.ai.*` model ids via the UC model-services API. + """List all `system.ai.*` model ids via the UC model-services API (see + :func:`list_model_service_entries`).""" + entries, reason = list_model_service_entries( + workspace, token, page_size=page_size, max_pages=max_pages, use_cache=use_cache + ) + return [entry.id for entry in entries], reason + + +def list_model_service_entries( + workspace: str, + token: str, + *, + page_size: int = _MODEL_SERVICES_PAGE_SIZE, + max_pages: int = 100, + use_cache: bool = True, +) -> tuple[list[ModelServiceEntry], str | None]: + """List all `system.ai.*` model services via the UC model-services API. Pages through ``/api/2.1/unity-catalog/model-services`` scoped to the ``system.ai`` schema (``parent=schemas/system.ai``) with a bounded - ``page_size`` (the endpoint 499s without one) and returns the de-duplicated, - sorted list of ``system.ai.`` ids. Returns (ids, reason); reason - is None on success, otherwise it describes why the list is empty (HTTP/network - error or no services). Scoping matters: the unscoped metastore listing walks - every schema across dozens of ~2s pages (~50s on a busy workspace) only to - keep the same ``system.ai.*`` subset — see ``_MODEL_SERVICE_PARENT_SCHEMA``. + ``page_size`` (the endpoint 499s without one) and returns the de-duplicated + entries sorted by id. Returns (entries, reason); reason is None on success, + otherwise it describes why the list is empty (HTTP/network error or no + services). Scoping matters: the unscoped metastore listing walks every schema + across dozens of ~2s pages (~50s on a busy workspace) only to keep the same + ``system.ai.*`` subset — see ``_MODEL_SERVICE_PARENT_SCHEMA``. A successful result is memoized per workspace for the life of the process; pass ``use_cache=False`` to force a fresh walk. @@ -2203,7 +2243,7 @@ def list_model_services( return list(cached), None origin = workspace_origin(workspace) - ids: list[str] = [] + entries: dict[str, ModelServiceEntry] = {} page_token: str | None = None seen_tokens: set[str] = set() last_reason: str | None = None @@ -2225,8 +2265,10 @@ def list_model_services( for service in data.get("model_services", []): if isinstance(service, dict): model_id = _model_service_id(service) - if model_id: - ids.append(model_id) + if model_id and model_id not in entries: + entries[model_id] = ModelServiceEntry( + model_id, _model_service_api_types(service) + ) page_token = data.get("next_page_token") or None if not page_token: last_reason = None @@ -2235,11 +2277,11 @@ def list_model_services( break seen_tokens.add(page_token) - deduped = sorted(set(ids)) - if deduped: + listed = [entries[model_id] for model_id in sorted(entries)] + if listed: if use_cache: - _MODEL_SERVICES_CACHE[workspace] = list(deduped) - return deduped, None + _MODEL_SERVICES_CACHE[workspace] = list(listed) + return listed, None return [], last_reason or "model-services listing returned no models" diff --git a/src/ucode/doctor.py b/src/ucode/doctor.py index 85ec1513a..c56566379 100644 --- a/src/ucode/doctor.py +++ b/src/ucode/doctor.py @@ -16,6 +16,7 @@ from ucode.agents import ( TOOL_SPECS, + copilot, tool_binary_installed, tool_version_error, update_tool_binary, @@ -184,6 +185,26 @@ def _check_agent_clis() -> list[Check]: return checks +def _check_copilot_models() -> Check | None: + """The models `ug copilot` registers for the current workspace, when Copilot is configured.""" + state = load_state() + if "copilot" not in (state.get("available_tools") or []): + return None + catalog = copilot.copilot_catalog(state) + if not catalog: + return Check( + "Copilot models", + "warn", + "none Copilot can use — run `ug configure --agents copilot` to rediscover them", + ) + counts = ", ".join( + f"{sum(m.provider == provider for m in catalog)} {provider}" + for provider in copilot.PROVIDER_NAMES.values() + if any(m.provider == provider for m in catalog) + ) + return Check("Copilot models", "ok", f"{len(catalog)} models ({counts})") + + def _check_databricks_auth() -> Check | None: """Validate the configured workspace's Databricks credentials. @@ -256,7 +277,7 @@ def _gather_checks() -> list[Check]: checks: list[Check] = [_check_uv(), _check_npm(), _check_databricks_cli(), _check_workspace()] # These return None when they don't apply (no workspace, no env collision), # so drop the Nones before display. - optional = [_check_databricks_auth(), _check_anthropic_env_collision()] + optional = [_check_databricks_auth(), _check_anthropic_env_collision(), _check_copilot_models()] checks.extend(c for c in optional if c is not None) checks.extend(_check_agent_clis()) checks.append(_check_ug()) diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 30d4f9ca5..0fe4ead21 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -98,7 +98,7 @@ def managed_state_overrides(managed: dict, tool: str) -> dict[str, object]: overrides["opencode_models"] = buckets else: overrides[f"{tool}_models"] = models - if tool in ("claude", "codex"): + if tool in ("claude", "codex", "copilot"): static_models = managed_static_models(managed, tool) if static_models: overrides[f"{tool}_static_models"] = static_models @@ -122,8 +122,13 @@ def managed_unservable_models(managed: dict, tool: str) -> list[str]: Only non-empty when *every* named model is unservable, which is when the translation yields nothing and the developer's own models stand — so the caller can say why the admin's list had no effect. opencode has no OpenAI provider and pi has no OSS provider, so each can be handed a - valid model FQN it cannot route. + valid model FQN it cannot route, and copilot serves only Claude and GPT. """ + if tool == "copilot": + from ucode.agents.copilot import classify_copilot_model + + listed = managed_static_models(managed, tool) or [] + return [] if any(classify_copilot_model(m) for m in listed) else listed if tool not in ("opencode", "pi"): return [] models = _manifest_models(managed, tool) diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index 8d130a057..a7feab5cb 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -58,7 +58,7 @@ "codex": ("codex", "oss"), "opencode": ("claude", "codex", "gemini", "oss"), "pi": ("claude", "codex", "gemini", "oss"), - "copilot": ("claude", "codex", "gemini", "oss"), + "copilot": ("claude", "codex"), } @@ -102,6 +102,11 @@ def model_options_for_agent(tool: str, state: dict) -> list[str]: models = state.get(key) if isinstance(models, list): options.extend(m for m in models if isinstance(m, str) and m) + if tool == "copilot": + from ucode.agents.copilot import classify_copilot_model + + # The codex listing also carries models Copilot can't serve, such as gpt-oss. + options = [m for m in options if classify_copilot_model(m)] # dict.fromkeys de-duplicates while preserving first-seen order (a model can match more than # one family bucket, e.g. an OSS id that also appears in the codex listing). return list(dict.fromkeys(options)) diff --git a/tests/AGENTS.md b/tests/AGENTS.md index 8ff486fbd..e6788b2e7 100644 --- a/tests/AGENTS.md +++ b/tests/AGENTS.md @@ -93,7 +93,7 @@ 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 (non-blocking lane). Do not add MCP/skills functionality, + model tasks and one TUI model-switch journey (non-blocking lane). Do not add 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 ba51b2caa..2f52f4c7e 100644 --- a/tests/README.md +++ b/tests/README.md @@ -432,7 +432,10 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_case_14_*` | Launch configured and fresh Codex with a model location | The app-server list exactly matches the independent API-compatible parent catalog, includes the dedicated Codex service, and contains no out-of-schema models | | `test_ug_claude_headless_prompt_argument`, `test_ug_claude_headless_prompt_stdin`, `test_ug_claude_headless_prompt_after_separator` | Run Claude from a script using each prompt form | Structured final answer contains the file value; exit zero; no routing | | `test_ug_codex_headless_prompt_argument`, `test_ug_codex_headless_prompt_stdin`, `test_ug_codex_headless_prompt_after_separator` | Run Codex from a script using each prompt form | Completed turn and final answer contain the file value; exit zero; no routing | -| `test_ug_copilot_headless_completes_file_task`, `test_ug_copilot_headless_gpt_responses_task` | Run Copilot CLI from a script (`-p ... --output-format json`) on a Claude model (Anthropic Messages) and a GPT model (Responses) | Completed turn whose assistant message contains the file value; no file under the session home holds the Databricks credential (non-blocking CI lane) | +| `test_ug_copilot_headless_completes_file_task`, `test_ug_copilot_headless_gpt_responses_task` | Run Copilot CLI from a script (`-p ... --output-format json`) on a Claude model (Anthropic Messages) and a GPT model (Responses) | Completed turn whose assistant message contains the file value and was sent as the requested model; the session started on the matching registry provider; no file under the session home holds the Databricks credential; no session registry left (non-blocking CI lane) | +| `test_ug_copilot_one_setup_runs_claude_then_gpt_by_provider_ref`, `test_ug_copilot_runs_an_older_claude_from_the_full_catalog` | One configure, then headless tasks on a Claude model, a GPT model chosen as `dbx-responses/`, and an older Claude from the full catalog | Each answer contains its file value and was sent as the requested model; sessions started on `dbx-anthropic` and `dbx-responses` (non-blocking CI lane) | +| `test_ug_copilot_unknown_model_lists_available_models` | Interactive launch with a model the workspace doesn't have | Exit 1 before Copilot starts; stderr names the model and lists the available ones, without `claude-opus-4-1`; no session registry left (non-blocking CI lane) | +| `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_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) | diff --git a/tests/integration/README.md b/tests/integration/README.md index 07e46b35c..5f6ddf469 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -211,6 +211,11 @@ Databricks credential. `test_ug_copilot_lifecycle.py` signs in through revoke them mid-task: once to require a refresh, once with sign-out to require the sign-in-expired message. It also covers configure/revert, including upgrading from ug v0.2.0, which saved the bearer. Refresh near natural expiry is covered by the hermetic proxy tests in `tests/test_gateway_proxy_integration.py`. +`test_ug_copilot_models.py` covers the per-session model registry: one setup runs a +Claude task and then a GPT task chosen with Copilot's `dbx-responses/` form, an +older Claude runs from the full catalog, an unknown model fails before Copilot starts, +and the TUI switches from Claude to GPT with `/model` mid-session. Model evidence comes +from Copilot's own `~/.copilot/session-state/*/events.jsonl` (`utils/agents/copilot.py`). ## Shared agent helpers @@ -426,6 +431,7 @@ test_ug_claude_tracing.py # Claude OTLP export reaches the configu test_ug_codex_headless.py # script prompts and model arguments test_ug_copilot_headless.py # Copilot via the session proxy: Claude + GPT tasks, no token on disk test_ug_copilot_lifecycle.py # Copilot refresh after revocation, sign-in expired, configure/revert, upgrade +test_ug_copilot_models.py # Copilot model registry: Claude then GPT, older Claude, unknown model, TUI /model switch test_ug_codex_tracing.py # Codex OTLP export reaches the configured trace table test_ug_claude_commands.py # command help forwarding test_ug_codex_commands.py # command help and parser error forwarding @@ -619,7 +625,7 @@ cases instead complete real file tasks through `--workspace` and `--model-locati journeys check setup and launch with existing services, without inference. There are **85 live cases** (including 11 marked TUI journeys) and **7 installation -checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds six live cases. One +checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds ten live cases. One **`workspace_switch` case** uses two real workspaces and checks skills MCP cleanup and a completed Claude task. A further **43 `managed_fixture` cases** (ten of them also `live`) run on the managed workspace with a checked-in JSON CodingAgentConfig from `tests/fixtures/managed_config/` @@ -845,7 +851,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`, six cases) runs on Linux and Windows the +A non-blocking **Copilot** job (`live and copilot`, ten cases) runs on Linux and Windows the same way and is likewise outside `cujs` until it is stable. Each lane installs only the agent CLIs it needs (one, except the managed OpenCode lane's diff --git a/tests/integration/test_ug_copilot_headless.py b/tests/integration/test_ug_copilot_headless.py index cae6f08a7..9fb0de7be 100644 --- a/tests/integration/test_ug_copilot_headless.py +++ b/tests/integration/test_ug_copilot_headless.py @@ -2,7 +2,13 @@ import pytest from utils.constants import CLAUDE_TEST_MODEL, CODEX_TEST_MODEL, COPILOT_HEADLESS_FLAGS -from utils.evidence import FileTask, assert_no_bearer_on_disk +from utils.evidence import ( + FileTask, + assert_completed_task_model, + assert_no_bearer_on_disk, + assert_no_session_registry_left, + copilot_selected_models, +) pytestmark = [pytest.mark.live, pytest.mark.copilot] @@ -10,9 +16,10 @@ def test_ug_copilot_headless_completes_file_task(live_session, workspace): """Scenario: configure Copilot, then run a headless Claude prompt through `ug copilot`. - Expected: Copilot reads the fixture through the session's loopback proxy (Anthropic - provider) and its completed answer contains the unknown value; no file under the - session home holds the Databricks credential afterwards. + Expected: Copilot starts on the registry's Anthropic provider, reads the fixture through + the session's loopback proxy, and its completed answer (sent as the requested model) + contains the unknown value; no file under the session home holds the Databricks + credential, and the session registry is gone afterwards. """ session = live_session task = FileTask(session) @@ -39,15 +46,19 @@ def test_ug_copilot_headless_completes_file_task(live_session, workspace): ) task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, CLAUDE_TEST_MODEL) + assert copilot_selected_models(session)[0] == f"dbx-anthropic/{CLAUDE_TEST_MODEL}" assert_no_bearer_on_disk(session) + assert_no_session_registry_left(session) def test_ug_copilot_headless_gpt_responses_task(live_session, workspace): """Scenario: configure Copilot, then run a headless GPT prompt through `ug copilot`. - Expected: Copilot reaches the gateway's Responses API through the session proxy and - its completed answer contains the fixture's unknown value; no file under the session - home holds the Databricks credential afterwards. + Expected: Copilot starts on the registry's Responses provider, reaches the gateway's + Responses API through the session proxy, and its completed answer (sent as the requested + model) contains the fixture's unknown value; no file under the session home holds the + Databricks credential, and the session registry is gone afterwards. """ session = live_session task = FileTask(session) @@ -74,4 +85,7 @@ def test_ug_copilot_headless_gpt_responses_task(live_session, workspace): ) task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, CODEX_TEST_MODEL) + assert copilot_selected_models(session)[0] == f"dbx-responses/{CODEX_TEST_MODEL}" assert_no_bearer_on_disk(session) + assert_no_session_registry_left(session) diff --git a/tests/integration/test_ug_copilot_models.py b/tests/integration/test_ug_copilot_models.py new file mode 100644 index 000000000..ef0145b74 --- /dev/null +++ b/tests/integration/test_ug_copilot_models.py @@ -0,0 +1,165 @@ +"""CUJs for discovering and switching between Databricks models in GitHub Copilot CLI through ug.""" + +import pytest +from utils.constants import ( + CLAUDE_TEST_MODEL, + CODEX_TEST_MODEL, + COPILOT_HEADLESS_FLAGS, + COPILOT_OLDER_CLAUDE_MODEL, +) +from utils.evidence import ( + FileTask, + assert_completed_task_model, + assert_no_session_registry_left, + copilot_selected_models, + fresh_id, +) +from utils.terminal import AgentTerminal + +pytestmark = [pytest.mark.live, pytest.mark.copilot] + + +def _configure(session, workspace): + session.run( + "configure", + "--agents", + "copilot", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + + +def test_ug_copilot_one_setup_runs_claude_then_gpt_by_provider_ref(live_session, workspace): + """Scenario: configure Copilot once, run a headless Claude task, then a headless GPT task + selected with Copilot's own `dbx-responses/` form, without reconfiguring. + + Expected: both completed answers contain their fixture values and were sent as the requested + models, the sessions started on `dbx-anthropic` then `dbx-responses`, and no session registry + is left afterwards. + """ + session = live_session + claude_task = FileTask(session) + gpt_task = FileTask(session) + _configure(session, workspace) + session.record("copilot-models.json", {"models": [CLAUDE_TEST_MODEL, CODEX_TEST_MODEL]}) + + claude = session.run( + "copilot", + "--model", + CLAUDE_TEST_MODEL, + "-p", + claude_task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + gpt = session.run( + "copilot", + "--model", + f"dbx-responses/{CODEX_TEST_MODEL}", + "-p", + gpt_task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + + claude_task.assert_headless_answer("copilot", claude) + gpt_task.assert_headless_answer("copilot", gpt) + assert_completed_task_model(session, "copilot", claude_task.value, CLAUDE_TEST_MODEL) + assert_completed_task_model(session, "copilot", gpt_task.value, CODEX_TEST_MODEL) + selected = copilot_selected_models(session) + assert f"dbx-anthropic/{CLAUDE_TEST_MODEL}" in selected + assert f"dbx-responses/{CODEX_TEST_MODEL}" in selected + assert_no_session_registry_left(session) + + +def test_ug_copilot_runs_an_older_claude_from_the_full_catalog(live_session, workspace): + """Scenario: configure Copilot, then run a headless task on an older Claude than its family's + newest model. + + Expected: the completed answer contains the fixture value and was sent as that older model, + which only reaches Copilot through the full discovered catalog. Limit: needs the model on the + e2e workspace. + """ + session = live_session + task = FileTask(session) + _configure(session, workspace) + session.record("copilot-model.json", {"model": COPILOT_OLDER_CLAUDE_MODEL}) + + result = session.run( + "copilot", + "--model", + COPILOT_OLDER_CLAUDE_MODEL, + "-p", + task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + + task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, COPILOT_OLDER_CLAUDE_MODEL) + + +def test_ug_copilot_unknown_model_lists_available_models(live_session, workspace): + """Scenario: configure Copilot, then launch it interactively with a model the workspace + doesn't have. + + Expected: ug exits 1 before starting Copilot, and its error names the model and lists the + available ones, including the Claude and GPT test models but not `claude-opus-4-1`, which + Copilot can't use. No session registry is left. + """ + session = live_session + _configure(session, workspace) + missing = f"system.ai.no-such-model-{fresh_id()}" + + result = session.run("copilot", "--model", missing, ok=False, timeout=120) + + assert result.returncode == 1, result.stdout + result.stderr + assert f"Model '{missing}' isn't available in ug copilot" in result.stderr + assert CLAUDE_TEST_MODEL in result.stderr + assert CODEX_TEST_MODEL in result.stderr + assert "claude-opus-4-1" not in result.stderr + assert not list(session.home.glob(".copilot/session-state/*")), "Copilot started" + assert_no_session_registry_left(session) + + +def test_ug_copilot_tui_switches_from_claude_to_gpt_with_model_command(live_session, workspace): + """Scenario: configure Copilot, start its TUI on a Claude model, complete a task, switch to a + GPT model with `/model dbx-responses/`, and complete a second task in the same session. + + Expected: the folder-trust prompt is answered through its visible choice, both answers carry + their fixture values and were sent as the Claude then the GPT model, the session records the + switch to the GPT ref, Copilot exits normally, and no session registry is left. + """ + session = live_session + claude_task = FileTask(session) + gpt_task = FileTask(session) + _configure(session, workspace) + session.record("copilot-models.json", {"models": [CLAUDE_TEST_MODEL, CODEX_TEST_MODEL]}) + gpt_ref = f"dbx-responses/{CODEX_TEST_MODEL}" + command = [ + str(session.binary), + "copilot", + "--model", + CLAUDE_TEST_MODEL, + "--allow-all-tools", + "--no-auto-update", + "--disable-builtin-mcps", + ] + + with AgentTerminal(session, "copilot", command, "copilot-model-switch") as tui: + tui.boot() + tui.submit(claude_task.prompt) + tui.wait_for_task(claude_task) + tui.submit(f"/model {gpt_ref}") + tui.wait_for(lambda screen: "Model changed" in screen, "the model switch notice") + tui.submit(gpt_task.prompt) + tui.wait_for_task(gpt_task) + tui.exit_normally() + + assert_completed_task_model(session, "copilot", claude_task.value, CLAUDE_TEST_MODEL) + assert_completed_task_model(session, "copilot", gpt_task.value, CODEX_TEST_MODEL) + assert copilot_selected_models(session)[-1] == gpt_ref + assert_no_session_registry_left(session) diff --git a/tests/integration/utils/agents/copilot.py b/tests/integration/utils/agents/copilot.py new file mode 100644 index 000000000..2ea73799d --- /dev/null +++ b/tests/integration/utils/agents/copilot.py @@ -0,0 +1,48 @@ +"""Copilot-specific integration evidence, read from Copilot's own session events.""" + +from . import _evidence_id + +__all__ = ["EVIDENCE_KIND", "assistant_answers", "completed_task_models", "selected_models"] + +SESSION_DIRECTORY = ".copilot/session-state" +EVIDENCE_KIND = "assistant message model" + + +def is_child_session(path: str, records: list[dict]) -> bool: + del path, records + return False + + +def _messages(records: list[dict]) -> list[dict]: + return [ + record.get("data") or {} for record in records if record.get("type") == "assistant.message" + ] + + +def assistant_answers(records: list[dict]) -> list[str]: + return [ + message["content"] + for message in _messages(records) + if isinstance(message.get("content"), str) + ] + + +def completed_task_models(records: list[dict], answer_value: str) -> set[str]: + """The model ids Copilot sent for the assistant messages that carry ``answer_value``.""" + return { + _evidence_id(message.get("model"), "model") + for message in _messages(records) + if answer_value in (message.get("content") or "") + } + + +def selected_models(records: list[dict]) -> list[str]: + """Registry refs the session started on or switched to, in order.""" + refs = [] + for record in records: + data = record.get("data") or {} + if record.get("type") == "session.start" and data.get("selectedModel"): + refs.append(data["selectedModel"]) + elif record.get("type") == "session.model_change" and data.get("newModel"): + refs.append(data["newModel"]) + return refs diff --git a/tests/integration/utils/constants.py b/tests/integration/utils/constants.py index a3adbd04c..2acc51a9f 100644 --- a/tests/integration/utils/constants.py +++ b/tests/integration/utils/constants.py @@ -6,6 +6,8 @@ # nano sometimes nests powershell inside Codex's pwsh, breaking `$` expansion in its file reads. CODEX_TEST_MODEL = "system.ai.gpt-5-6-luna" if os.name == "nt" else "system.ai.gpt-5-4-nano" +# An older Claude than the family default: only the full catalog registers it for Copilot. +COPILOT_OLDER_CLAUDE_MODEL = "system.ai.claude-sonnet-4-5" COPILOT_HEADLESS_FLAGS = ( "--allow-all-tools", "--output-format", diff --git a/tests/integration/utils/evidence.py b/tests/integration/utils/evidence.py index 0e6e32c55..4dfbf1d8c 100644 --- a/tests/integration/utils/evidence.py +++ b/tests/integration/utils/evidence.py @@ -5,9 +5,9 @@ import uuid from pathlib import Path -from .agents import _evidence_id, claude, codex +from .agents import _evidence_id, claude, codex, copilot -_AGENT_HELPERS = {"claude": claude, "codex": codex} +_AGENT_HELPERS = {"claude": claude, "codex": codex, "copilot": copilot} def assert_no_terminal_api_error(screen: str) -> None: @@ -147,6 +147,21 @@ def assert_no_bearer_on_disk(session) -> None: assert "OAUTH_TOKEN" not in text +def copilot_selected_models(session) -> list[str]: + """Every registry ref Copilot sessions under this home started on or switched to.""" + return [ + ref + for records in agent_sessions(session, "copilot").values() + for ref in copilot.selected_models(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(".ucode/run/copilot-*")) + assert not leftovers, f"Session registries left behind: {leftovers}" + + def fresh_id(): """Return a 12-hex-char id; models drop characters when copying longer ones.""" return uuid.uuid4().hex[:12] diff --git a/tests/integration/utils/terminal.py b/tests/integration/utils/terminal.py index 9a35c9e29..64da87814 100644 --- a/tests/integration/utils/terminal.py +++ b/tests/integration/utils/terminal.py @@ -229,7 +229,7 @@ def __exit__(self, exc_type, exc, traceback): "normal_exit_observed": self.ended and self.child.exitstatus == 0, "routing_log": routing_log.read_text() if routing_log.is_file() else None, "agent_sessions": agent_sessions(self.session, self.agent) - if self.agent in ("claude", "codex") + if self.agent in ("claude", "codex", "copilot") else {}, }, ) @@ -403,6 +403,14 @@ def boot(self, timeout=120): and bool(re.search(r"1[.)]\s+Yes, I trust (?:this|the) folder", text)), "\r", ), + ( + "copilot-folder-trust", + self.agent == "copilot" + and "Confirm folder trust" in text + and self.session.cwd.name in text + and bool(re.search(r"❯\s*1[.)]\s+Yes\b", text)), + "\r", + ), ( "trust-directory", self.agent == "codex" @@ -444,9 +452,10 @@ def boot(self, timeout=120): "Configured ug launched Codex's account-login flow instead of its gateway session:\n" + text ) - title = "Claude Code" if self.agent == "claude" else "Codex" + title = {"claude": "Claude Code", "copilot": "Copilot v"}.get(self.agent, "Codex") if ( title in text + and "Confirm folder trust" not in text and "loading" not in text.lower() and not re.search(r"(?m)^\s*>\s+You are in\b", text) and re.search(r"(?m)^\s*[❯›>]\s*(?!\d+[.)])", text) diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 71e995ec4..3739d1308 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -3,14 +3,23 @@ from __future__ import annotations import json +import os import httpx import pytest from ucode import custom_oauth from ucode.agents import copilot +from ucode.databricks import ModelServiceEntry WS = "https://example.databricks.com" +PROXY = "http://127.0.0.1:4321" +SONNET = "system.ai.claude-sonnet-5-5" +GPT = "system.ai.gpt-5-4" + + +def catalog(*ids: str) -> list[copilot.CopilotModel]: + return copilot.build_copilot_catalog(ids) class TestCopilotSpec: @@ -27,111 +36,382 @@ def test_config_path_is_ucode_env_file(self): assert copilot.SPEC["config_path"].name == "ucode.env" -class TestRenderEnvOverlay: - def test_sets_provider_base_url(self): - env = copilot.render_env_overlay(WS, "gpt-5") - assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" +class TestClassifyCopilotModel: + @pytest.mark.parametrize( + ("model", "wire"), + [ + ("system.ai.claude-sonnet-5-5", "anthropic"), + ("system.ai.claude-sonnet-5-5[1m]", "anthropic"), + ("databricks-claude-sonnet-4-6", "anthropic"), + ("main.models.claude-opus-4-7", "anthropic"), + ("system.ai.gpt-5-4", "responses"), + ("system.ai.gpt-6-1-sol", "responses"), + ("databricks-gpt-5-3-codex", "responses"), + ("system.ai.gpt-4-1", "completions"), + ], + ) + def test_names_decide_the_wire(self, model, wire): + assert copilot.classify_copilot_model(model) == wire - def test_claude_uses_anthropic_provider_with_catalog_model_id(self): - env = copilot.render_env_overlay(WS, "system.ai.claude-sonnet-5-5") + @pytest.mark.parametrize( + ("model", "reason"), + [ + ("system.ai.gemini-3-flash", "Gemini"), + ("system.ai.gpt-oss-120b", "Claude and GPT"), + ("system.ai.meta-llama-3-3-70b-instruct", "Claude and GPT"), + ("system.ai.gte-large-en", "Claude and GPT"), + ("system.ai.claude-opus-4-1", "thinking"), + ("anthropic-aigw-12e8d1ee-system.ai.gpt-5-6-sol", "alias"), + ], + ) + def test_excluded_models_say_why(self, model, reason): + assert copilot.classify_copilot_model(model) is None + assert reason in (copilot.copilot_exclusion_reason(model) or "") + + def test_advertised_routes_only_exclude(self): + assert copilot.classify_copilot_model(SONNET, ["mlflow/v1/embeddings"]) is None + # gpt-5-3-codex lists only openai/codex Responses, which the gateway also serves on mlflow. + codex_routes = ["openai/v1/responses", "codex/v1/responses"] + assert ( + copilot.classify_copilot_model("system.ai.gpt-5-3-codex", codex_routes) == "responses" + ) - assert env["COPILOT_PROVIDER_TYPE"] == "anthropic" - assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/anthropic" - assert env["COPILOT_PROVIDER_MODEL_ID"] == "claude-sonnet-5.5" - assert env["COPILOT_MODEL"] == "system.ai.claude-sonnet-5-5" - assert "COPILOT_PROVIDER_WIRE_API" not in env - def test_claude_wire_model_override_selects_anthropic_provider(self): - env = copilot.render_env_overlay(WS, "gpt-5", override_model="system.ai.claude-opus-4-7") +class TestBuildCopilotCatalog: + def test_orders_claude_by_family_newest_first_then_gpt(self): + models = catalog( + "system.ai.gpt-5-4", + "system.ai.claude-haiku-4-5", + "system.ai.claude-sonnet-4-6", + "system.ai.gpt-6-1-sol", + "system.ai.claude-sonnet-5-5", + "system.ai.claude-opus-5-5", + "system.ai.claude-fable-5-1", + "system.ai.gemini-3-flash", + ) - assert env["COPILOT_PROVIDER_TYPE"] == "anthropic" - assert env["COPILOT_PROVIDER_MODEL_ID"] == "claude-opus-4.7" + assert [m.id for m in models] == [ + "system.ai.claude-fable-5-1", + "system.ai.claude-opus-5-5", + "system.ai.claude-sonnet-5-5", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5", + "system.ai.gpt-6-1-sol", + "system.ai.gpt-5-4", + ] + + def test_uses_the_listing_routes_and_drops_duplicates(self): + models = copilot.build_copilot_catalog( + [ + ModelServiceEntry(SONNET, ("anthropic/v1/messages",)), + ModelServiceEntry("system.ai.gte-large-en", ("mlflow/v1/embeddings",)), + SONNET, + ] + ) - def test_sets_provider_type(self): - env = copilot.render_env_overlay(WS, "m") - assert env["COPILOT_PROVIDER_TYPE"] == "openai" + assert models == [copilot.CopilotModel(SONNET, "anthropic")] + assert models[0].ref == f"dbx-anthropic/{SONNET}" - def test_sets_model(self): - env = copilot.render_env_overlay(WS, "claude-sonnet-4-6") - assert env["COPILOT_MODEL"] == "claude-sonnet-4-6" - @pytest.mark.parametrize("model", ["m", "system.ai.claude-sonnet-5-5"]) - def test_carries_no_credential(self, model): - env = copilot.render_env_overlay(WS, model) - for key in copilot.CREDENTIAL_ENV_KEYS: - assert key not in env +class TestCatalogFromListings: + def test_prefers_the_model_services_listing(self): + entries = [ModelServiceEntry(GPT), ModelServiceEntry("system.ai.gemini-3-flash")] + assert copilot.catalog_from_listings(entries, [SONNET]) == catalog(GPT) - def test_sets_offline_true(self): - env = copilot.render_env_overlay(WS, "m") - assert env["COPILOT_OFFLINE"] == "true" + def test_falls_back_to_the_gateway_listings(self): + assert copilot.catalog_from_listings(None, ["databricks-claude-sonnet-4-6"]) == catalog( + "databricks-claude-sonnet-4-6" + ) + + def test_no_answer_keeps_the_saved_catalog(self): + assert copilot.catalog_from_listings(None, []) is None + + def test_an_answer_with_nothing_usable_is_an_empty_catalog(self): + assert ( + copilot.catalog_from_listings([ModelServiceEntry("system.ai.gemini-3-flash")], []) == [] + ) + + +class TestCopilotCatalog: + def test_managed_list_wins_and_skips_what_copilot_cannot_use(self): + state = { + "copilot_static_models": [GPT, "system.ai.gemini-3-flash"], + "copilot_catalog": copilot.catalog_state(catalog(SONNET)), + } + + assert copilot.copilot_catalog(state) == catalog(GPT) + assert copilot.unservable_static_models(state) == [ + ( + "system.ai.gemini-3-flash", + copilot.copilot_exclusion_reason("system.ai.gemini-3-flash"), + ) + ] + + def test_an_unusable_managed_list_falls_back_to_discovery(self): + state = { + "copilot_static_models": ["system.ai.gemini-3-flash"], + "copilot_catalog": copilot.catalog_state(catalog(SONNET)), + } + + assert copilot.copilot_catalog(state) == catalog(SONNET) + assert copilot.unservable_static_models(state) == [] + + def test_saved_catalog(self): + saved = copilot.catalog_state(catalog(SONNET, GPT)) + + assert saved["models"] == [SONNET, GPT] + assert saved["discovered_at"].endswith("+00:00") + assert copilot.copilot_catalog({"copilot_catalog": saved}) == catalog(SONNET, GPT) + + def test_an_empty_saved_catalog_does_not_revive_older_lists(self): + state = {"copilot_catalog": copilot.catalog_state([]), "codex_models": [GPT]} + assert copilot.copilot_catalog(state) == [] + + def test_state_from_before_the_catalog_uses_the_claude_and_gpt_lists(self): + state = { + "claude_models": {"sonnet": SONNET}, + "codex_models": [GPT, "system.ai.gpt-oss-120b"], + } + assert copilot.copilot_catalog(state) == catalog(SONNET, GPT) + + +class TestResolveCopilotModel: + MODELS = catalog(SONNET, GPT, "databricks-claude-opus-4-7") @pytest.mark.parametrize( - ("selected_model", "override_model", "expected_api"), + "requested", [ - ("gpt-6.1-sol", None, "responses"), - ("system.ai.gpt-6-astra", "gpt-4.1", "completions"), - ("system.ai.gpt-5-6-sol", "databricks-gpt-6-1-sol", "responses"), - ("system.ai.gpt-6-astra", "", "responses"), - ("gpt-5", "", "responses"), - ("gpt-4.1", "", "completions"), + SONNET, + "claude-sonnet-5-5", + "CLAUDE-Sonnet-5-5", + f"{SONNET}[1m]", + f"dbx-anthropic/{SONNET}", ], ) - def test_selects_wire_api_from_override_and_keeps_selected_model( - self, selected_model, override_model, expected_api - ): - env = copilot.render_env_overlay(WS, selected_model, override_model=override_model) + def test_accepts_ids_short_names_and_refs(self, requested): + assert copilot.resolve_copilot_model(requested, self.MODELS).id == SONNET - assert env["COPILOT_PROVIDER_WIRE_API"] == expected_api - assert env["COPILOT_MODEL"] == selected_model + def test_accepts_the_databricks_prefix(self): + assert copilot.resolve_copilot_model("claude-opus-4-7", self.MODELS).id == ( + "databricks-claude-opus-4-7" + ) + def test_prefers_the_system_ai_id(self): + models = catalog("databricks-claude-sonnet-5-5", SONNET) + assert copilot.resolve_copilot_model("claude-sonnet-5-5", models).id == SONNET -PROXY = "http://127.0.0.1:4321" + def test_unknown_model_lists_what_is_available(self): + with pytest.raises(RuntimeError) as error: + copilot.resolve_copilot_model("system.ai.gemini-3-flash", self.MODELS) + message = str(error.value) + assert "Gemini" in message + assert f"dbx-anthropic: databricks-claude-opus-4-7, {SONNET}" in message + assert f"dbx-responses: {GPT}" in message -class TestBuildRuntimeEnv: - def test_inherits_path(self): - env = copilot.build_runtime_env(PROXY, "m", "secret") - assert "PATH" in env - def test_points_copilot_at_the_session_proxy(self): - env = copilot.build_runtime_env(PROXY, "m", "secret") - assert env["COPILOT_PROVIDER_BASE_URL"] == f"{PROXY}/ai-gateway/mlflow/v1" - assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "secret" +class TestLaunchModel: + MODELS = catalog(SONNET, "system.ai.claude-opus-5-5", GPT) + + def test_requested_model_wins(self): + state = {"copilot_default_model": SONNET} + assert copilot.launch_model(state, self.MODELS, requested=GPT).id == GPT + + def test_managed_default_then_list(self): + assert copilot.launch_model({"copilot_default_model": GPT}, self.MODELS).id == GPT + assert copilot.launch_model({"copilot_static_models": [GPT]}, catalog(GPT)).id == GPT + + @pytest.mark.parametrize( + ("ids", "expected"), + [ + ((GPT, "system.ai.claude-opus-5-5", SONNET), SONNET), + ((GPT, "system.ai.claude-opus-5-5"), "system.ai.claude-opus-5-5"), + (("system.ai.claude-fable-5", GPT), GPT), + ], + ) + def test_defaults_to_sonnet_then_opus_then_gpt(self, ids, expected): + assert copilot.launch_model({}, catalog(*ids)).id == expected + + def test_an_unavailable_requested_model_warns_and_falls_through(self, capsys): + model = copilot.launch_model({}, self.MODELS, requested="system.ai.gemini-3-flash") + + assert model.id == SONNET + assert "gemini-3-flash isn't available in ug copilot" in capsys.readouterr().out - def test_claude_goes_through_the_proxy_anthropic_route(self): - env = copilot.build_runtime_env(PROXY, "system.ai.claude-sonnet-5-5", "secret") - assert env["COPILOT_PROVIDER_BASE_URL"] == f"{PROXY}/ai-gateway/anthropic" + def test_default_model_is_quiet(self, capsys): + assert copilot.default_model({"copilot_default_model": "missing"}) is None + state = {"copilot_default_model": "missing", "codex_models": [GPT]} + assert copilot.default_model(state) == GPT + assert capsys.readouterr().out == "" - def test_sets_no_oauth_token(self, monkeypatch): - monkeypatch.delenv("OAUTH_TOKEN", raising=False) - env = copilot.build_runtime_env(PROXY, "m", "secret") - assert "OAUTH_TOKEN" not in env + +class TestRewriteModelArgs: + MODELS = catalog(SONNET, GPT) + + @pytest.mark.parametrize( + "args", + [["--model", GPT], [f"--model={GPT}"], ["-m", "gpt-5-4"], ["--model", SONNET, "-m", GPT]], + ) + def test_rewrites_to_the_registry_ref(self, args): + rewritten, selected = copilot.rewrite_model_args(["-p", "hi", *args], self.MODELS) + + assert selected == copilot.CopilotModel(GPT, "responses") + assert rewritten[:2] == ["-p", "hi"] + assert rewritten[-2:] == ["--model", f"dbx-responses/{GPT}"] + + def test_leaves_arguments_after_the_separator(self): + rewritten, selected = copilot.rewrite_model_args(["--", "--model", "x"], self.MODELS) + assert (rewritten, selected) == (["--", "--model", "x"], None) + + def test_unknown_model_raises(self): + with pytest.raises(RuntimeError, match="isn't available"): + copilot.rewrite_model_args(["--model", "system.ai.claude-opus-4-1"], self.MODELS) + + +class TestCopilotCatalogModelId: + @pytest.mark.parametrize( + ("model", "expected"), + [ + ("system.ai.claude-sonnet-5-5", "claude-sonnet-5.5"), + ("system.ai.claude-opus-4-7[1m]", "claude-opus-4.7"), + ("system.ai.claude-sonnet-5", "claude-sonnet-5"), + ("system.ai.claude-haiku-4-5", "claude-haiku-4.5"), + ("system.ai.claude-haiku-5-5", "claude-sonnet-5.5"), + ("databricks-claude-sonnet-4-6", "claude-sonnet-4.6"), + ], + ) + def test_maps_claude_ids(self, model, expected): + assert copilot.copilot_catalog_model_id(model) == expected + + def test_other_models_have_no_catalog_id(self): + assert copilot.copilot_catalog_model_id(GPT) is None + + +class TestCopilotModelLimits: + @pytest.mark.parametrize( + ("model", "limits"), + [ + ("system.ai.claude-sonnet-5-5", (872_000, 128_000)), + ("system.ai.claude-haiku-5-5", (872_000, 128_000)), + ("system.ai.claude-haiku-4-5", (136_000, 64_000)), + ("system.ai.claude-sonnet-4-5", (168_000, 32_000)), + ("system.ai.claude-opus-4-8", (936_000, 32_000)), + ("system.ai.claude-opus-6", (936_000, 32_000)), + ("system.ai.claude-sonnet-4-7", (936_000, 32_000)), + ("system.ai.claude-opus-4-0", (168_000, 32_000)), + ], + ) + def test_claude_limits(self, model, limits): + assert copilot.copilot_model_limits(model) == limits + + def test_gpt_keeps_copilots_defaults(self): + assert copilot.copilot_model_limits(GPT) is None + + +class TestRenderProvidersConfig: + def test_one_provider_per_wire_in_use_on_the_proxy(self): + config = copilot.render_providers_config(catalog(SONNET, GPT), PROXY, "secret") + + assert config["providers"] == [ + { + "name": "dbx-anthropic", + "type": "anthropic", + "baseUrl": f"{PROXY}/ai-gateway/anthropic", + "bearerToken": "secret", + }, + { + "name": "dbx-responses", + "type": "openai", + "wireApi": "responses", + "baseUrl": f"{PROXY}/ai-gateway/mlflow/v1", + "bearerToken": "secret", + }, + ] + + def test_completions_provider(self): + config = copilot.render_providers_config(catalog("system.ai.gpt-4-1"), PROXY, "secret") + assert [(p["name"], p.get("wireApi")) for p in config["providers"]] == [ + ("dbx-completions", "completions") + ] + + def test_models_carry_claude_settings_and_limits(self): + config = copilot.render_providers_config(catalog(SONNET, GPT), PROXY, "secret") + + assert config["models"] == [ + { + "id": SONNET, + "provider": "dbx-anthropic", + "wireModel": SONNET, + "name": SONNET, + "modelId": "claude-sonnet-5.5", + "maxPromptTokens": 872_000, + "maxOutputTokens": 128_000, + }, + {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": GPT}, + ] + + +class TestBuildRuntimeEnv: + def test_points_copilot_at_the_registry(self, tmp_path): + env = copilot.build_runtime_env(tmp_path / "providers.json", catalog(GPT)[0]) + + assert env["COPILOT_PROVIDERS_CONFIG"] == str(tmp_path / "providers.json") + assert env["COPILOT_MODEL"] == f"dbx-responses/{GPT}" + assert env["COPILOT_OFFLINE"] == "true" + assert env["PATH"] == os.environ["PATH"] @pytest.mark.parametrize( "key", [ - "COPILOT_PROVIDER_API_KEY", + "COPILOT_PROVIDERS_CONFIG", + "COPILOT_PROVIDER_BASE_URL", + "COPILOT_PROVIDER_BEARER_TOKEN", "COPILOT_PROVIDER_API_KEY_COMMAND", - "COPILOT_PROVIDER_TRANSPORT", + "COPILOT_PROVIDER_WIRE_MODEL", ], ) - def test_drops_inherited_settings_that_bypass_the_proxy(self, monkeypatch, key): + def test_drops_the_users_own_provider_settings(self, monkeypatch, tmp_path, key): monkeypatch.setenv(key, "user-value") - env = copilot.build_runtime_env(PROXY, "m", "secret") - assert key not in env + env = copilot.build_runtime_env(tmp_path / "providers.json", catalog(GPT)[0]) + assert env.get(key) != "user-value" - def test_inherited_wire_model_overrides_route_without_being_cleared(self, monkeypatch): - monkeypatch.setenv("COPILOT_PROVIDER_WIRE_MODEL", "gpt-4.1") - monkeypatch.setenv("COPILOT_PROVIDER_WIRE_API", "responses") - monkeypatch.setenv("COPILOT_PROVIDER_MODEL_ID", "user-model-id") - monkeypatch.setenv("COPILOT_PROVIDER_MODEL_LIMITS_ID", "user-model-limits") - env = copilot.build_runtime_env(PROXY, "system.ai.gpt-6-astra", "secret") +class TestSessionRegistry: + def test_writes_a_private_registry_removed_at_exit(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + + with copilot.session_registry({"models": ["x"]}) as path: + assert json.loads(path.read_text(encoding="utf-8")) == {"models": ["x"]} + assert path.parent.parent == copilot.COPILOT_RUN_DIR + if os.name != "nt": + assert path.stat().st_mode & 0o777 == 0o600 + assert path.parent.stat().st_mode & 0o777 == 0o700 + + assert not path.parent.exists() - assert env["COPILOT_PROVIDER_WIRE_MODEL"] == "gpt-4.1" - assert env["COPILOT_PROVIDER_WIRE_API"] == "completions" - assert env["COPILOT_PROVIDER_MODEL_ID"] == "user-model-id" - assert env["COPILOT_PROVIDER_MODEL_LIMITS_ID"] == "user-model-limits" + def test_removed_when_the_session_fails(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + + with pytest.raises(OSError), copilot.session_registry({}) as path: + raise OSError("copilot failed") + + assert not path.parent.exists() + + def test_sweep_keeps_live_sessions_and_removes_dead_ones(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + dead = copilot.COPILOT_RUN_DIR / "copilot-dead" + dead.mkdir(parents=True) + (dead / "session.lock").touch() + (dead / "providers.json").write_text("{}", encoding="utf-8") + unlocked = copilot.COPILOT_RUN_DIR / "copilot-crashed-before-lock" + unlocked.mkdir() + + with copilot.session_registry({}) as live: + copilot.sweep_session_dirs() + assert live.exists() + + assert not dead.exists() + assert not unlocked.exists() class TestMcpServerConfig: @@ -230,95 +510,6 @@ def test_removes_mcp_server_without_clobbering_others(self, tmp_path, monkeypatc assert written["other"] is True -class TestDefaultModel: - def test_prefers_claude_sonnet(self): - state = { - "claude_models": {"sonnet": "s4", "opus": "o4", "haiku": "h4"}, - "codex_models": ["gpt-5"], - } - assert copilot.default_model(state) == "s4" - - def test_falls_back_to_opus(self): - state = {"claude_models": {"opus": "o4", "haiku": "h4"}} - assert copilot.default_model(state) == "o4" - - def test_falls_back_to_haiku(self): - state = {"claude_models": {"haiku": "h4"}} - assert copilot.default_model(state) == "h4" - - def test_falls_back_to_codex_when_no_claude(self): - state = {"codex_models": ["gpt-5", "gpt-5-mini"]} - assert copilot.default_model(state) == "gpt-5" - - def test_returns_none_when_no_models(self): - assert copilot.default_model({}) is None - - def test_ignores_gemini_models(self): - # Gemini is excluded — Databricks' Gemini translator rejects copilot's request shape. - state = {"gemini_models": ["gemini-2-5-pro"]} - assert copilot.default_model(state) is None - - -class TestCopilotCatalogModelId: - @pytest.mark.parametrize( - ("model", "expected"), - [ - ("system.ai.claude-sonnet-5-5", "claude-sonnet-5.5"), - ("system.ai.claude-opus-5-5[1m]", "claude-opus-5.5"), - ("system.ai.claude-opus-4-7", "claude-opus-4.7"), - ("system.ai.claude-sonnet-5", "claude-sonnet-5"), - ("system.ai.claude-haiku-4-5", "claude-haiku-4.5"), - ("system.ai.claude-haiku-5-5", "claude-sonnet-5.5"), - ("system.ai.claude-haiku-6", "claude-sonnet-6"), - ], - ) - def test_maps_claude_ids(self, model, expected): - assert copilot.copilot_catalog_model_id(model) == expected - - @pytest.mark.parametrize( - "model", - [ - "gpt-5", - "system.ai.gpt-6-astra", - "claude-sonnet-4-6", - "databricks-claude-opus-4-7", - "global.anthropic.claude-opus-4-8", - "main.schema.my-claude-finetune", - "system.ai.myclaude-model", - ], - ) - def test_non_system_claude_ids_return_none(self, model): - assert copilot.copilot_catalog_model_id(model) is None - - -class TestModelUsesResponsesApi: - def test_numeric_gpt_majors_and_gateway_aliases_use_responses(self): - for model in ( - "gpt-5", - "databricks-gpt-5-mini", - "system.ai.gpt-5-4", - "gpt-6", - "gpt-6.1-sol", - "system.ai.gpt-6-astra", - "databricks-gpt-6-1-sol", - "gpt-7", - "gpt-10", - ): - assert copilot.model_uses_responses_api(model), model - - def test_pre_gpt5_non_gpt_and_unsupported_alias_shapes_use_completions(self): - for model in ( - "gpt-4.1", - "gpt-4o", - "claude-sonnet-4-6", - "my-gpt-6-model", - "gpt6", - "gpt.6", - "gpt-60x", - ): - assert not copilot.model_uses_responses_api(model), model - - def isolate_copilot_config_paths(monkeypatch, tmp_path): import ucode.config_io as config_io import ucode.state as state_mod @@ -331,61 +522,32 @@ def isolate_copilot_config_paths(monkeypatch, tmp_path): 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_RUN_DIR", app_dir / "run") return env_path class TestWriteToolConfig: - def test_manages_wire_api_and_preserves_wire_model_and_user_values(self, tmp_path, monkeypatch): + def test_records_the_default_model_and_drops_stale_provider_settings( + self, tmp_path, monkeypatch + ): env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) env_path.parent.mkdir(parents=True) env_path.write_text( - "COPILOT_PROVIDER_WIRE_MODEL=system.ai.gpt-6-1-sol\n" - "COPILOT_PROVIDER_WIRE_API=completions\n" - "COPILOT_PROVIDER_MODEL_ID=user-model-id\n" - "COPILOT_PROVIDER_MODEL_LIMITS_ID=user-model-limits\n" - "USER_SETTING=keep\n", + 'COPILOT_PROVIDER_TYPE="anthropic"\nCOPILOT_PROVIDER_WIRE_API="responses"\n' + 'USER_SETTING="kept"\n', encoding="utf-8", ) - monkeypatch.setenv("COPILOT_PROVIDER_WIRE_MODEL", "gpt-5") - - state = copilot.write_tool_config({"workspace": WS}, "system.ai.gpt-5-6-sol") - - written = copilot.parse_dotenv(env_path) - assert written["COPILOT_MODEL"] == "system.ai.gpt-5-6-sol" - assert written["COPILOT_PROVIDER_WIRE_API"] == "responses" - assert written["COPILOT_PROVIDER_WIRE_MODEL"] == "system.ai.gpt-6-1-sol" - assert "COPILOT_PROVIDER_MODEL_ID" not in written - assert written["COPILOT_PROVIDER_MODEL_LIMITS_ID"] == "user-model-limits" - assert written["USER_SETTING"] == "keep" - assert "COPILOT_PROVIDER_WIRE_API" in state["managed_configs"]["copilot"]["keys"] - assert "COPILOT_PROVIDER_WIRE_MODEL" not in state["managed_configs"]["copilot"]["keys"] - - def test_claude_drops_stale_wire_api_from_env_file(self, tmp_path, monkeypatch): - env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) - env_path.parent.mkdir(parents=True) - env_path.write_text("COPILOT_PROVIDER_WIRE_API=completions\n", encoding="utf-8") - - copilot.write_tool_config({"workspace": WS}, "system.ai.claude-sonnet-5-5") + state = {"workspace": WS, "codex_models": [GPT]} - written = copilot.parse_dotenv(env_path) - assert "COPILOT_PROVIDER_WIRE_API" not in written - assert written["COPILOT_PROVIDER_TYPE"] == "anthropic" - assert written["COPILOT_PROVIDER_MODEL_ID"] == "claude-sonnet-5.5" - - def test_needs_no_token_and_writes_none(self, tmp_path, monkeypatch): - env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) + state = copilot.write_tool_config(state, GPT) - def no_token(*args, **kwargs): - raise AssertionError("configure must not mint a token") - - monkeypatch.setattr(custom_oauth, "get_databricks_token", no_token) - - copilot.write_tool_config({"workspace": WS}, "gpt-5") - - written = copilot.parse_dotenv(env_path) - assert written["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" - for key in copilot.CREDENTIAL_ENV_KEYS: - assert key not in written + env = copilot.parse_dotenv(env_path) + assert env == { + "USER_SETTING": "kept", + "COPILOT_MODEL": f"dbx-responses/{GPT}", + "COPILOT_OFFLINE": "true", + } + assert state["managed_configs"]["copilot"]["keys"] == copilot.MANAGED_KEYS def test_scrubs_a_token_an_earlier_version_wrote(self, tmp_path, monkeypatch): env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) @@ -531,65 +693,91 @@ def _stub_launch(monkeypatch, process=None, token="dbx-token"): def popen(argv, *, env): seen["argv"] = argv seen["env"] = env - base = env["COPILOT_PROVIDER_BASE_URL"] - seen["accepted"] = httpx.post( - f"{base}/chat/completions", - headers={"Authorization": "Bearer wrong-secret"}, - content=b"{}", - timeout=5, - ).status_code + registry_path = env["COPILOT_PROVIDERS_CONFIG"] + seen["registry_path"] = registry_path + seen["registry"] = json.loads(open(registry_path, encoding="utf-8").read()) + provider = seen["registry"]["providers"][0] + seen["accepted"] = { + secret: httpx.post( + f"{provider['baseUrl']}/v1/messages", + headers={"Authorization": f"Bearer {secret}"}, + content=b"{}", + timeout=5, + ).status_code + for secret in ("wrong-secret",) + } return process or _ExitedProcess() monkeypatch.setattr(copilot.subprocess_cross_os, "popen", popen) return seen +STATE = {"workspace": WS, "copilot_catalog": {"models": [SONNET, GPT]}} + + class TestLaunch: - @pytest.mark.parametrize( - ("tool_args", "pinned_model", "default", "expected_model", "expected_api"), - [ - (["--model", "gpt-6.1-sol"], "gpt-5", "gpt-5", "gpt-6.1-sol", "responses"), - ([], "system.ai.gpt-4-1", "gpt-6", "system.ai.gpt-4-1", "completions"), - ([], None, "system.ai.gpt-10", "system.ai.gpt-10", "responses"), - ], - ) - def test_launch_model_precedence_preserves_argv( - self, tool_args, pinned_model, default, expected_model, expected_api, tmp_path, monkeypatch - ): + def test_copilot_gets_the_registry_and_never_the_token(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) - monkeypatch.delenv("COPILOT_PROVIDER_WIRE_MODEL", raising=False) - seen = _stub_launch(monkeypatch) - state = {"workspace": WS, "copilot_default_model": default} + seen = _stub_launch(monkeypatch, token="dbx-token") with pytest.raises(SystemExit) as exit_info: - copilot.launch( - state, - tool_args, - options=copilot.LaunchOptions(user_pinned_model=pinned_model), - ) + copilot.launch(dict(STATE), [], options=copilot.LaunchOptions()) assert exit_info.value.code == 0 - assert seen["argv"] == ["copilot", *tool_args] - assert seen["env"]["COPILOT_MODEL"] == expected_model - assert seen["env"]["COPILOT_PROVIDER_WIRE_API"] == expected_api - assert state["copilot_default_model"] == default + assert seen["env"]["COPILOT_MODEL"] == f"dbx-anthropic/{SONNET}" + assert "dbx-token" not in seen["env"].values() + registry = seen["registry"] + assert [m["id"] for m in registry["models"]] == [SONNET, GPT] + secret = registry["providers"][0]["bearerToken"] + assert secret not in ("", "dbx-token") + assert secret not in seen["env"].values() + assert "dbx-token" not in json.dumps(registry) + # The proxy is live while Copilot runs and refuses other credentials. + assert seen["accepted"] == {"wrong-secret": 401} + assert not os.path.exists(seen["registry_path"]) - def test_copilot_gets_the_proxy_and_a_session_secret_not_the_token(self, tmp_path, monkeypatch): + def test_explicit_model_is_rewritten_to_the_registry_ref(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) - monkeypatch.delenv("COPILOT_PROVIDER_WIRE_MODEL", raising=False) - seen = _stub_launch(monkeypatch, token="dbx-token") + seen = _stub_launch(monkeypatch) with pytest.raises(SystemExit): copilot.launch( - {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") + dict(STATE), + ["-p", "hi", "-m", "gpt-5-4"], + options=copilot.LaunchOptions(launch_model=SONNET), ) - env = seen["env"] - assert env["COPILOT_PROVIDER_BASE_URL"].startswith("http://127.0.0.1:") - assert env["COPILOT_PROVIDER_BEARER_TOKEN"] not in ("", "dbx-token") - assert "dbx-token" not in env.values() - # The proxy is live while Copilot runs and refuses other credentials. - assert seen["accepted"] == 401 + assert seen["argv"] == ["copilot", "-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): + isolate_copilot_config_paths(monkeypatch, tmp_path) + seen = _stub_launch(monkeypatch) + + with pytest.raises(SystemExit): + copilot.launch(dict(STATE), [], options=copilot.LaunchOptions(launch_model=GPT)) + + assert seen["env"]["COPILOT_MODEL"] == f"dbx-responses/{GPT}" + + def test_unknown_model_fails_before_starting_copilot(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + seen = _stub_launch(monkeypatch) + + with pytest.raises(RuntimeError, match="isn't available in ug copilot"): + copilot.launch(dict(STATE), ["--model", "nope"], options=copilot.LaunchOptions()) + + assert seen == {} + assert not copilot.COPILOT_RUN_DIR.exists() + + def test_warns_about_managed_models_copilot_cannot_use(self, tmp_path, monkeypatch, capsys): + isolate_copilot_config_paths(monkeypatch, tmp_path) + _stub_launch(monkeypatch) + state = {"workspace": WS, "copilot_static_models": [GPT, "system.ai.gemini-3-flash"]} + + with pytest.raises(SystemExit): + copilot.launch(state, [], options=copilot.LaunchOptions()) + + assert "Skipping system.ai.gemini-3-flash" in capsys.readouterr().out def test_launch_persists_no_state_and_no_token(self, tmp_path, monkeypatch): env_path = isolate_copilot_config_paths(monkeypatch, tmp_path) @@ -598,9 +786,7 @@ def test_launch_persists_no_state_and_no_token(self, tmp_path, monkeypatch): monkeypatch.setattr(copilot, "save_state", saved.append) with pytest.raises(SystemExit): - copilot.launch( - {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") - ) + copilot.launch(dict(STATE), [], options=copilot.LaunchOptions()) assert saved == [] assert not env_path.exists() @@ -610,31 +796,30 @@ def test_ctrl_c_keeps_waiting_for_copilot(self, tmp_path, monkeypatch): _stub_launch(monkeypatch, process=_ExitedProcess(returncode=130, interrupts=2)) with pytest.raises(SystemExit) as exit_info: - copilot.launch( - {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") - ) + copilot.launch(dict(STATE), [], options=copilot.LaunchOptions()) assert exit_info.value.code == 130 - def test_proxy_stops_when_copilot_cannot_start(self, tmp_path, monkeypatch): + def test_proxy_and_registry_stop_when_copilot_cannot_start(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) monkeypatch.setattr(custom_oauth, "get_databricks_token", lambda *args, **kwargs: "tok") monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") seen = {} def popen(argv, *, env): - seen["base"] = env["COPILOT_PROVIDER_BASE_URL"] + seen["registry"] = env["COPILOT_PROVIDERS_CONFIG"] + registry = json.loads(open(seen["registry"], encoding="utf-8").read()) + seen["base"] = registry["providers"][0]["baseUrl"] raise FileNotFoundError("copilot") monkeypatch.setattr(copilot.subprocess_cross_os, "popen", popen) with pytest.raises(FileNotFoundError): - copilot.launch( - {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") - ) + copilot.launch(dict(STATE), [], options=copilot.LaunchOptions()) + assert not os.path.exists(seen["registry"]) with pytest.raises(httpx.ConnectError): - httpx.post(f"{seen['base']}/responses", content=b"{}", timeout=10) + httpx.post(f"{seen['base']}/v1/messages", content=b"{}", timeout=10) def test_no_model_is_a_launch_error(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) @@ -644,21 +829,34 @@ def test_no_model_is_a_launch_error(self, tmp_path, monkeypatch): class TestManagedKeys: - def test_includes_required_vars(self): - for key in ( - "COPILOT_PROVIDER_TYPE", - "COPILOT_PROVIDER_BASE_URL", - "COPILOT_PROVIDER_WIRE_API", - "COPILOT_MODEL", - "COPILOT_OFFLINE", - ): - assert key in copilot.MANAGED_KEYS + def test_only_the_model_and_offline_mode(self): + assert copilot.MANAGED_KEYS == ["COPILOT_MODEL", "COPILOT_OFFLINE"] def test_excludes_credentials(self): for key in copilot.CREDENTIAL_ENV_KEYS: assert key not in copilot.MANAGED_KEYS +class TestMinimumVersion: + @pytest.mark.parametrize( + ("version", "blocked"), + [ + ("GitHub Copilot CLI 1.0.90.", True), + ("1.0.91", False), + ("1.1.0", False), + ("unknown", False), + ], + ) + def test_requires_1_0_91(self, monkeypatch, version, blocked): + monkeypatch.setattr(copilot, "agent_version", lambda binary: version) + + error = copilot.minimum_version_error() + + assert (error is not None) == blocked + if blocked: + assert "npm install -g @github/copilot" in error + + class TestValidateCmd: def test_starts_with_binary(self): cmd = copilot.validate_cmd("copilot") @@ -678,25 +876,6 @@ def test_adds_ucode_mcp_config_when_present(self, tmp_path, monkeypatch): assert cmd[:3] == ["copilot", "--additional-mcp-config", f"@{mcp_path}"] -class TestManagedModels: - def test_managed_models_win_over_the_shared_discovery_lists(self): - state = { - "copilot_models": ["system.ai.gpt-5"], - "claude_models": {"sonnet": "shared-should-not-win"}, - } - assert copilot.default_model(state) == "system.ai.gpt-5" - - def test_falls_back_to_the_shared_lists_without_a_managed_config(self): - assert copilot.default_model({"claude_models": {"sonnet": "discovered"}}) == "discovered" - - def test_copilot_default_model_wins_over_allowlist(self): - state = { - "copilot_default_model": "admin-chosen-default", - "copilot_models": ["system.ai.gpt-5"], - } - assert copilot.default_model(state) == "admin-chosen-default" - - class TestWriteUserMcpServers: def test_batched_add_remove_preserves_other_keys(self, tmp_path, monkeypatch): path = tmp_path / "ucode-mcp-config.json" diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 47b4acacb..b6a94b1ed 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -287,10 +287,11 @@ def test_opencode_available(self): assert check_gateway_endpoint(state, "opencode") is True def test_copilot_available_with_claude(self): - assert check_gateway_endpoint({"claude_models": {"sonnet": "s4"}}, "copilot") is True + state = {"claude_models": {"sonnet": "system.ai.claude-sonnet-4-6"}} + assert check_gateway_endpoint(state, "copilot") is True def test_copilot_available_with_codex(self): - assert check_gateway_endpoint({"codex_models": ["m"]}, "copilot") is True + assert check_gateway_endpoint({"codex_models": ["system.ai.gpt-5-4"]}, "copilot") is True def test_copilot_unavailable_with_only_gemini(self): # Gemini is intentionally excluded from Copilot. @@ -1090,9 +1091,9 @@ def test_fable_only_workspace_has_a_default(tool): def test_fable_does_not_displace_existing_default(tool): state = { "claude_models": {"fable": "system.ai.claude-fable-5"}, - "codex_models": ["existing-gpt"], + "codex_models": ["system.ai.gpt-5-4"], } - assert default_model_for_tool(tool, state) == "existing-gpt" + assert default_model_for_tool(tool, state) == "system.ai.gpt-5-4" class TestConfigureSelectedTools: diff --git a/tests/test_cli.py b/tests/test_cli.py index 9462d2236..7da9229bf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3469,6 +3469,24 @@ def test_reports_runtime_error(self): assert "boom" in _strip_ansi(result.output) +class TestCopilotLaunchModel: + MANAGED = { + "enabled_agents": {"copilot": {"model_config": {"default_model": "system.ai.gpt-5-4"}}} + } + + def test_managed_default_is_the_launch_model(self): + with _launch_policy_patches(self.MANAGED) as calls: + result = runner.invoke(app, ["copilot"]) + + assert result.exit_code == 0, result.output + assert calls["launch"].call_args.kwargs["options"].launch_model == "system.ai.gpt-5-4" + + def test_a_managed_default_still_discovers_the_catalog(self): + assert not cli_mod._skip_model_discovery( + self.MANAGED, "copilot", provider=None, managed_parent_schema=None, family_defaults=None + ) + + class TestOpenCodeModelFlag: @pytest.mark.parametrize("flag", ["--model", "-m"]) def test_model_option_threads_to_config_and_launch(self, flag): @@ -5675,6 +5693,56 @@ def _boom(*a, **k): assert state["claude_models"] == {"opus": "databricks-claude-opus-4-8"} +class TestConfigureSharedStateCopilotCatalog: + """Copilot's catalog comes from the same listings as the other agents' models.""" + + WS = "https://copilot.databricks.com" + + def _configure(self, monkeypatch, *, listed, saved=None): + import ucode.cli as cli_mod + import ucode.databricks as db_mod + + TestConfigureSharedStateSkipDiscovery._stub(monkeypatch) + previous = {"workspace": self.WS} + if saved is not None: + previous["copilot_catalog"] = saved + monkeypatch.setattr(cli_mod, "load_state", lambda: dict(previous)) + if listed is not None: + db_mod._MODEL_SERVICES_CACHE[self.WS] = [db_mod.ModelServiceEntry(m) for m in listed] + reason = None if listed else "HTTP 503" + sonnet = "system.ai.claude-sonnet-4-6" + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ( + ({"sonnet": sonnet}, [], [], [], None) if listed else ({}, [], [], [], reason) + ), + ) + monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, reason)) + monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], reason)) + return cli_mod.configure_shared_state(self.WS, tools=["copilot"]) + + def test_saves_the_classified_catalog(self, monkeypatch): + state = self._configure( + monkeypatch, + listed=["system.ai.claude-sonnet-4-6", "system.ai.gpt-5-4", "system.ai.gemini-3-flash"], + ) + + assert state["copilot_catalog"]["models"] == [ + "system.ai.claude-sonnet-4-6", + "system.ai.gpt-5-4", + ] + + def test_failed_listing_keeps_the_saved_catalog(self, monkeypatch, capsys): + saved = {"models": ["system.ai.gpt-5-4"], "discovered_at": "2026-10-01T00:00:00+00:00"} + + state = self._configure(monkeypatch, listed=None, saved=saved) + + assert state["copilot_catalog"] == saved + assert state["_discovery_reasons"]["copilot"] == "HTTP 503" + assert "Copilot model discovery failed (HTTP 503)" in capsys.readouterr().out + + class TestConfigureSharedStateSkipPreflight: """With skip_preflight (--skip-preflight), a prior configure is trusted: no auth login, token fetch, gateway probe, or model discovery runs — but the @@ -6492,6 +6560,19 @@ def test_other_agent_commands_keep_stdout(self): assert cli_mod._child_owns_stdout("claude", ["app-server"]) is False assert cli_mod._child_owns_stdout("gemini", []) is False + @pytest.mark.parametrize( + ("tool_args", "expected"), + [ + (["-p", "hi"], True), + (["--prompt", "hi"], True), + (["--prompt=hi"], True), + ([], False), + (["--", "-p", "hi"], False), + ], + ) + def test_copilot_prompt_mode_owns_stdout(self, tool_args, expected): + assert cli_mod._child_owns_stdout("copilot", tool_args) is expected + @pytest.mark.parametrize( "option", [ diff --git a/tests/test_databricks.py b/tests/test_databricks.py index d2b569ea6..9910d1a1f 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -707,6 +707,34 @@ def flaky_get(url, token, timeout=10): assert ids == ["system.ai.gpt-5"] assert calls["n"] == 3 # two failures, third succeeds + def test_entries_keep_the_advertised_api_types(self, monkeypatch): + service = { + **_model_service("system.ai.gte-large-en"), + "supported_api_types": ["mlflow/v1/embeddings", 3], + } + + monkeypatch.setattr( + db_mod, + "_get_model_services_page", + lambda url, token: ({"model_services": [service]}, None), + ) + + entries, reason = db_mod.list_model_service_entries(WS, "token") + + assert reason is None + assert entries == [ + db_mod.ModelServiceEntry("system.ai.gte-large-en", ("mlflow/v1/embeddings",)) + ] + assert db_mod.cached_model_service_entries(WS) == entries + + def test_failed_listing_is_not_cached(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_get_model_services_page", lambda url, token: (None, "HTTP 503") + ) + + assert db_mod.list_model_service_entries(WS, "token") == ([], "HTTP 503") + assert db_mod.cached_model_service_entries(WS) is None + class TestModelServiceExists: def test_true_when_listed_in_its_schema(self, monkeypatch): diff --git a/tests/test_doctor.py b/tests/test_doctor.py index 45ea62c40..31e68fc39 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -143,6 +143,33 @@ def test_unknown_tool_is_skipped(self): assert _check_agent_clis() == [] +class TestCopilotModelsCheck: + def test_skipped_when_copilot_is_not_configured(self): + with patch.object(doctor_mod, "load_state", return_value={"available_tools": ["claude"]}): + assert doctor_mod._check_copilot_models() is None + + def test_counts_models_per_provider(self): + state = { + "available_tools": ["copilot"], + "copilot_catalog": {"models": ["system.ai.claude-sonnet-4-6", "system.ai.gpt-5-4"]}, + } + with patch.object(doctor_mod, "load_state", return_value=state): + check = doctor_mod._check_copilot_models() + + assert check is not None + assert check.status == "ok" + assert check.detail == "2 models (1 dbx-anthropic, 1 dbx-responses)" + + def test_warns_when_nothing_is_usable(self): + state = {"available_tools": ["copilot"], "copilot_catalog": {"models": []}} + with patch.object(doctor_mod, "load_state", return_value=state): + check = doctor_mod._check_copilot_models() + + assert check is not None + assert check.status == "warn" + assert "ug configure --agents copilot" in check.detail + + class TestDatabricksAuthCheck: def test_none_when_no_workspace(self): with patch.object(doctor_mod, "load_state", return_value={}): diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 92d86a216..09f20aa81 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -1041,7 +1041,7 @@ def test_recovers_from_rejected_initial_token( class TestCopilotLaunch: - """Run Copilot through MLflow using Responses for GPT-6+ and Chat Completions otherwise. + """Run Copilot on every catalog model through one session registry. Gemini is excluded by design — Databricks' Gemini translator rejects the `stream_options` field Copilot CLI sends. Other incompatible models are @@ -1095,10 +1095,8 @@ def test_launch_copilot_per_model( pytest.skip("No Copilot-compatible models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - env_path = tmp_path / ".copilot-env" - backup_path = tmp_path / "copilot-env.backup" - monkeypatch.setattr(copilot, "COPILOT_ENV_PATH", env_path) - monkeypatch.setattr(copilot, "COPILOT_BACKUP_PATH", backup_path) + monkeypatch.setattr(copilot, "COPILOT_RUN_DIR", tmp_path / "run") + catalog = copilot.build_copilot_catalog([model for _, model in models]) from ucode import gateway_proxy @@ -1116,21 +1114,19 @@ def test_launch_copilot_per_model( proxy_base_url = f"http://127.0.0.1:{server.server_address[1]}" failures = [] + config = copilot.render_providers_config(catalog, proxy_base_url, session_secret) try: - for family, model in models: - with pytest.MonkeyPatch().context() as mp: - mp.setattr("ucode.agents.copilot.save_state", lambda s: None) - copilot.write_tool_config({**e2e_state, "workspace": e2e_workspace}, model) - - env = copilot.build_runtime_env(proxy_base_url, model, session_secret) - cmd = copilot.validate_cmd("copilot") - result = _run_agent(cmd, env=env, timeout=120) - combined = (result.stdout + result.stderr).strip() - if result.returncode != 0 or not combined: - failures.append( - f"family={family} model={model} rc={result.returncode} " - f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" - ) + with copilot.session_registry(config) as registry_path: + for model in catalog: + env = copilot.build_runtime_env(registry_path, model) + cmd = [*copilot.validate_cmd("copilot"), "--model", model.ref] + result = _run_agent(cmd, env=env, timeout=120) + combined = (result.stdout + result.stderr).strip() + if result.returncode != 0 or not combined: + failures.append( + f"model={model.ref} rc={result.returncode} " + f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" + ) finally: cache.stop() server.shutdown() diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index 4bf2b5ab5..f3125411b 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -683,6 +683,22 @@ def test_no_warning_when_anything_is_servable(self, tool, models): def test_agents_that_pass_models_through_never_warn(self): assert managed_unservable_models(self._managed("codex", ["anything"]), "codex") == [] + @staticmethod + def _copilot(models): + return {"enabled_agents": {"copilot": {"model_config": {"model_services": models}}}} + + def test_copilot_list_with_only_gemini_and_oss_is_unservable(self): + models = ["system.ai.gemini-3-flash", "system.ai.kimi-k2-6"] + assert managed_unservable_models(self._copilot(models), "copilot") == models + + def test_copilot_list_with_a_claude_model_is_servable(self): + models = ["system.ai.gemini-3-flash", "system.ai.claude-opus-4-8"] + assert managed_unservable_models(self._copilot(models), "copilot") == [] + + def test_copilot_model_list_reaches_state_as_its_static_models(self): + overrides = managed_state_overrides(self._copilot(["system.ai.gpt-5-4"]), "copilot") + assert overrides["copilot_static_models"] == ["system.ai.gpt-5-4"] + class TestRecommendedAgent: """A tier can move the org to a cheaper agent; with none named, default_agent stands.""" diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index e06260915..e4c9be432 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -333,13 +333,20 @@ def test_codex_sees_gpt_and_oss(self): ] def test_multi_provider_agents_see_everything(self): - for tool in ("opencode", "pi", "copilot"): + for tool in ("opencode", "pi"): options = model_options_for_agent(tool, STATE) assert "system.ai.claude-opus-4-8" in options, tool assert "system.ai.gpt-5-6" in options, tool assert "system.ai.gemini-3-flash" in options, tool assert "system.ai.kimi-k2-6" in options, tool + def test_copilot_sees_only_claude_and_gpt(self): + options = model_options_for_agent("copilot", STATE) + assert "system.ai.claude-opus-4-8" in options + assert "system.ai.gpt-5-6" in options + assert "system.ai.gemini-3-flash" not in options + assert "system.ai.kimi-k2-6" not in options + def test_empty_state_yields_no_options(self): assert model_options_for_agent("claude", {}) == [] From 46fd13dbefe1e177064ceb72cd66458312439540 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Tue, 20 Oct 2026 12:18:15 +0000 Subject: [PATCH 08/11] copilot: keep the session registry in ~/.copilot/ucode-providers-.json Co-authored-by: Isaac --- README.md | 4 +- src/ucode/agents/copilot.py | 60 ++++++++++++++--------------- src/ucode/cli.py | 2 +- tests/integration/utils/evidence.py | 2 +- tests/test_agent_copilot.py | 32 ++++++++------- tests/test_e2e.py | 2 +- 6 files changed, 53 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index deccce8ee..e7a1cd550 100644 --- a/README.md +++ b/README.md @@ -82,8 +82,8 @@ global choice is saved for plain `copilot` too, where the Databricks models don' exist. Requires Copilot CLI 1.0.91 or newer. `ug copilot` runs Copilot against a loopback proxy for the session. Copilot holds only -a random session secret, kept in the session registry under `~/.ucode/run/` (readable -only by you, deleted when the session ends); the proxy adds a fresh Databricks token to +a random session secret, kept in the session's `~/.copilot/ucode-providers-.json` +registry (readable only by you, deleted when the session ends); the proxy adds a fresh Databricks token to each gateway request, so OAuth sessions keep working past token expiry without a restart, and no Databricks token is written to disk. If the Databricks sign-in itself expires, Copilot shows a "Databricks sign-in expired" error; run `databricks auth login` and diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index d3f443da3..51138fcb0 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -1,7 +1,8 @@ """GitHub Copilot CLI agent: runs Copilot against a per-session Databricks model registry. Copilot reads its BYOK models from a providers.json registry (`COPILOT_PROVIDERS_CONFIG`). Each -`ug copilot` launch writes one under ~/.ucode/run/ listing every Databricks model Copilot can use, +`ug copilot` launch writes one (~/.copilot/ucode-providers-.json) listing every Databricks +model Copilot can use, so `/model` switches between them mid-session, across wire APIs. Each family keeps its native wire: Claude on the gateway's Anthropic endpoint (with Copilot's catalog id as `modelId`, so Copilot sends Claude request settings), GPT 5+ on Responses (Copilot sends `temperature: 0` on Chat @@ -22,8 +23,6 @@ import os import re import secrets -import shutil -import tempfile import threading import time from collections.abc import Callable, Iterable, Iterator @@ -31,7 +30,7 @@ from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path -from typing import IO, Any, Protocol +from typing import Any, Protocol from ucode.config_io import ( APP_DIR, @@ -70,11 +69,8 @@ COPILOT_BACKUP_PATH = APP_DIR / "copilot-ucode-env.backup" COPILOT_MCP_BACKUP_PATH = APP_DIR / "copilot-ucode-mcp-config.backup.json" COPILOT_PROXY_LOG_PATH = APP_DIR / "logs" / "copilot-proxy.log" -COPILOT_RUN_DIR = APP_DIR / "run" -_SESSION_DIR_PREFIX = "copilot-" -_SESSION_LOCK_NAME = "session.lock" -_RUN_DIR_LOCK_NAME = ".copilot.lock" -_REGISTRY_NAME = "providers.json" +_REGISTRY_PREFIX = "ucode-providers-" +_REGISTRY_LOCK_NAME = "ucode-providers.lock" MINIMUM_COPILOT_VERSION = (1, 0, 91) @@ -480,9 +476,9 @@ def build_runtime_env(registry_path: Path, model: CopilotModel) -> dict[str, str @contextmanager -def _run_dir_lock() -> Iterator[None]: - COPILOT_RUN_DIR.mkdir(parents=True, exist_ok=True) - with (COPILOT_RUN_DIR / _RUN_DIR_LOCK_NAME).open("a+b") as lock_file: +def _registries_lock() -> Iterator[None]: + COPILOT_CONFIG_DIR.mkdir(parents=True, exist_ok=True) + with (COPILOT_CONFIG_DIR / _REGISTRY_LOCK_NAME).open("a+b") as lock_file: acquire_exclusive_file_lock(lock_file) try: yield @@ -490,9 +486,9 @@ def _run_dir_lock() -> Iterator[None]: release_file_lock(lock_file) -def _sweep_dead_sessions() -> None: - for session_dir in COPILOT_RUN_DIR.glob(f"{_SESSION_DIR_PREFIX}*"): - lock_path = session_dir / _SESSION_LOCK_NAME +def _sweep_dead_registries() -> None: + for registry in COPILOT_CONFIG_DIR.glob(f"{_REGISTRY_PREFIX}*.json"): + lock_path = registry.with_suffix(".lock") if lock_path.exists(): with lock_path.open("a+b") as lock_file: try: @@ -500,44 +496,46 @@ def _sweep_dead_sessions() -> None: except (TimeoutError, OSError): continue release_file_lock(lock_file) - shutil.rmtree(session_dir, ignore_errors=True) + registry.unlink(missing_ok=True) + lock_path.unlink(missing_ok=True) -def sweep_session_dirs() -> None: +def sweep_session_registries() -> None: """Delete registries left by sessions that ended without cleaning up (e.g. a killed terminal).""" - if not COPILOT_RUN_DIR.exists(): + if not COPILOT_CONFIG_DIR.exists(): return - with _run_dir_lock(): - _sweep_dead_sessions() + with _registries_lock(): + _sweep_dead_registries() @contextmanager def session_registry(config: dict) -> Iterator[Path]: """Write ``config`` to a private per-session providers.json, removed when the session ends. - The session holds a lock on its directory so a concurrent sweep leaves it alone.""" - lock_file: IO[Any] | None = None - with _run_dir_lock(): - _sweep_dead_sessions() - session_dir = Path(tempfile.mkdtemp(prefix=_SESSION_DIR_PREFIX, dir=COPILOT_RUN_DIR)) + The session holds a lock on a sibling ``.lock`` file so a concurrent sweep leaves it alone; + locking the registry itself would block Copilot from reading it on Windows.""" + with _registries_lock(): + _sweep_dead_registries() + path = COPILOT_CONFIG_DIR / f"{_REGISTRY_PREFIX}{secrets.token_hex(8)}.json" + lock_path = path.with_suffix(".lock") + lock_file = lock_path.open("a+b") try: - lock_file = (session_dir / _SESSION_LOCK_NAME).open("a+b") acquire_exclusive_file_lock(lock_file, timeout=0) - path = session_dir / _REGISTRY_NAME descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) with os.fdopen(descriptor, "w", encoding="utf-8") as registry: json.dump(config, registry, indent=2) except BaseException: - if lock_file is not None: - lock_file.close() - shutil.rmtree(session_dir, ignore_errors=True) + lock_file.close() + path.unlink(missing_ok=True) + lock_path.unlink(missing_ok=True) raise try: yield path finally: release_file_lock(lock_file) lock_file.close() - shutil.rmtree(session_dir, ignore_errors=True) + path.unlink(missing_ok=True) + lock_path.unlink(missing_ok=True) def _parse_version(value: str) -> tuple[int, int, int] | None: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a3882f542..ae9bab563 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1421,7 +1421,7 @@ def revert() -> int: mcp_results = revert_mcp_configs(state) # Before restoring backups: one written by an earlier ug may hold a token. copilot_agent.scrub_stored_credentials() - copilot_agent.sweep_session_dirs() + copilot_agent.sweep_session_registries() claude_managed_result = claude_agent.revert_managed_settings() codex_managed_result = codex_agent.revert_managed_config() diff --git a/tests/integration/utils/evidence.py b/tests/integration/utils/evidence.py index 4dfbf1d8c..e2975c2fe 100644 --- a/tests/integration/utils/evidence.py +++ b/tests/integration/utils/evidence.py @@ -158,7 +158,7 @@ def copilot_selected_models(session) -> list[str]: 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(".ucode/run/copilot-*")) + leftovers = sorted(session.home.glob(".copilot/ucode-providers-*")) assert not leftovers, f"Session registries left behind: {leftovers}" diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 3739d1308..c101fa121 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -382,12 +382,13 @@ def test_writes_a_private_registry_removed_at_exit(self, tmp_path, monkeypatch): with copilot.session_registry({"models": ["x"]}) as path: assert json.loads(path.read_text(encoding="utf-8")) == {"models": ["x"]} - assert path.parent.parent == copilot.COPILOT_RUN_DIR + assert path.parent == copilot.COPILOT_CONFIG_DIR + assert path.name.startswith("ucode-providers-") if os.name != "nt": assert path.stat().st_mode & 0o777 == 0o600 - assert path.parent.stat().st_mode & 0o777 == 0o700 - assert not path.parent.exists() + assert not path.exists() + assert not path.with_suffix(".lock").exists() def test_removed_when_the_session_fails(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) @@ -395,23 +396,28 @@ def test_removed_when_the_session_fails(self, tmp_path, monkeypatch): with pytest.raises(OSError), copilot.session_registry({}) as path: raise OSError("copilot failed") - assert not path.parent.exists() + assert not path.exists() def test_sweep_keeps_live_sessions_and_removes_dead_ones(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) - dead = copilot.COPILOT_RUN_DIR / "copilot-dead" - dead.mkdir(parents=True) - (dead / "session.lock").touch() - (dead / "providers.json").write_text("{}", encoding="utf-8") - unlocked = copilot.COPILOT_RUN_DIR / "copilot-crashed-before-lock" - unlocked.mkdir() + config_dir = copilot.COPILOT_CONFIG_DIR + config_dir.mkdir(parents=True) + dead = config_dir / "ucode-providers-dead.json" + dead.write_text("{}", encoding="utf-8") + dead.with_suffix(".lock").touch() + unlocked = config_dir / "ucode-providers-crashed-before-lock.json" + unlocked.write_text("{}", encoding="utf-8") + users_own = config_dir / "providers.json" + users_own.write_text("{}", encoding="utf-8") with copilot.session_registry({}) as live: - copilot.sweep_session_dirs() + copilot.sweep_session_registries() assert live.exists() assert not dead.exists() + assert not dead.with_suffix(".lock").exists() assert not unlocked.exists() + assert users_own.exists() class TestMcpServerConfig: @@ -522,7 +528,7 @@ def isolate_copilot_config_paths(monkeypatch, tmp_path): 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_RUN_DIR", app_dir / "run") + monkeypatch.setattr(copilot, "COPILOT_CONFIG_DIR", env_path.parent) return env_path @@ -767,7 +773,7 @@ def test_unknown_model_fails_before_starting_copilot(self, tmp_path, monkeypatch copilot.launch(dict(STATE), ["--model", "nope"], options=copilot.LaunchOptions()) assert seen == {} - assert not copilot.COPILOT_RUN_DIR.exists() + assert not list(copilot.COPILOT_CONFIG_DIR.glob("ucode-providers-*")) def test_warns_about_managed_models_copilot_cannot_use(self, tmp_path, monkeypatch, capsys): isolate_copilot_config_paths(monkeypatch, tmp_path) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 09f20aa81..03353873f 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -1095,7 +1095,7 @@ def test_launch_copilot_per_model( pytest.skip("No Copilot-compatible models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - monkeypatch.setattr(copilot, "COPILOT_RUN_DIR", tmp_path / "run") + monkeypatch.setattr(copilot, "COPILOT_CONFIG_DIR", tmp_path / ".copilot") catalog = copilot.build_copilot_catalog([model for _, model in models]) from ucode import gateway_proxy From 5d4bc4afc8247f772810f73d8528a467539f024a Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Tue, 20 Oct 2026 13:39:48 +0000 Subject: [PATCH 09/11] copilot: readable model names in the session registry Co-authored-by: Isaac --- src/ucode/agents/claude.py | 29 +++++------------------------ src/ucode/agents/copilot.py | 3 ++- src/ucode/databricks.py | 23 +++++++++++++++++++++++ tests/test_agent_copilot.py | 4 ++-- 4 files changed, 32 insertions(+), 27 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 251d92fb7..4b229a9f9 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -48,6 +48,7 @@ build_tool_base_url, get_databricks_token, maybe_add_gpt_1m_suffix, + model_display_label, ug_binary, ) from ucode.launcher import exec_or_spawn @@ -602,7 +603,7 @@ def render_overlay( overlay["modelPicker"] = { "replaceBuiltInOptions": True, "options": [ - {"model": maybe_add_gpt_1m_suffix(m), "label": _picker_label(m)} + {"model": maybe_add_gpt_1m_suffix(m), "label": model_display_label(m)} for m in static_models ], } @@ -613,7 +614,8 @@ def render_overlay( "options": [ _picker_option( maybe_add_gpt_1m_suffix(model), - picker_catalog.model_id_to_display_name.get(model) or _picker_label(model), + picker_catalog.model_id_to_display_name.get(model) + or model_display_label(model), picker_catalog.model_id_to_description.get(model), ) for model in picker_catalog.model_ids @@ -632,32 +634,11 @@ def render_overlay( return overlay, keys -_MODEL_LABEL_ACRONYMS = frozenset({"glm", "gpt"}) _CLAUDE_PICKER_LABEL_RE = re.compile( r"^Claude (Fable|Opus|Sonnet|Haiku) (\d+(?:\.\d+)*)$", re.IGNORECASE ) -def _picker_label(model: str) -> str: - """A human-friendly picker label for a model id (e.g. ``system.ai.claude-haiku-4-5`` -> - ``Claude Haiku 4.5``): keep the vendor and name words title-cased, uppercase known acronyms, - and join a run of numeric segments into a dotted version.""" - stem = model.removeprefix("system.ai.") - parts: list[str] = [] - version: list[str] = [] - for token in stem.split("-"): - if token.isdigit(): - version.append(token) - continue - if version: - parts.append(".".join(version)) - version = [] - parts.append(token.upper() if token in _MODEL_LABEL_ACRONYMS else token.title()) - if version: - parts.append(".".join(version)) - return " ".join(parts) if parts else stem - - def _picker_option(model: str, label: str, description: str | None = None) -> dict[str, str]: option = {"model": model, "label": label} if description: @@ -700,7 +681,7 @@ def default_model_picker_catalog( descriptions: dict[str, str] = {} for family, raw_model in defaults.items(): model = raw_model - label = _picker_label(model.removesuffix("[1m]")) + label = model_display_label(model.removesuffix("[1m]")) if provider is not None: # Family shortcuts stay distinct from catalog rows for the same target. model = family diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 51138fcb0..236505e4e 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -49,6 +49,7 @@ ModelServiceEntry, build_copilot_base_url, build_tool_base_url, + model_display_label, model_version_sort_key, ) from ucode.os_compatibility import subprocess_cross_os @@ -448,7 +449,7 @@ def render_providers_config( "id": model.id, "provider": model.provider, "wireModel": model.id, - "name": model.id, + "name": model_display_label(model.id), } catalog_id = copilot_catalog_model_id(model.id) if model.wire == "anthropic" else None if catalog_id is not None: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 8c5321764..157f9110a 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3512,6 +3512,29 @@ def fetch_ai_gateway_claude_models(workspace: str, token: str) -> dict[str, str] return models +_MODEL_LABEL_ACRONYMS = frozenset({"glm", "gpt"}) + + +def model_display_label(model: str) -> str: + """A human-friendly picker label for a model id (e.g. ``system.ai.claude-haiku-4-5`` -> + ``Claude Haiku 4.5``): keep the vendor and name words title-cased, uppercase known acronyms, + and join a run of numeric segments into a dotted version.""" + stem = model.removeprefix("system.ai.") + parts: list[str] = [] + version: list[str] = [] + for token in stem.split("-"): + if token.isdigit(): + version.append(token) + continue + if version: + parts.append(".".join(version)) + version = [] + parts.append(token.upper() if token in _MODEL_LABEL_ACRONYMS else token.title()) + if version: + parts.append(".".join(version)) + return " ".join(parts) if parts else stem + + def model_version_sort_key(name: str) -> tuple: """Sort endpoint names so newer model versions come first. diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index c101fa121..248f5135c 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -342,12 +342,12 @@ def test_models_carry_claude_settings_and_limits(self): "id": SONNET, "provider": "dbx-anthropic", "wireModel": SONNET, - "name": SONNET, + "name": "Claude Sonnet 5.5", "modelId": "claude-sonnet-5.5", "maxPromptTokens": 872_000, "maxOutputTokens": 128_000, }, - {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": GPT}, + {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": "GPT 5.4"}, ] From 70b8f0e0cb4fad32dff715812070c5616aca74aa Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 21 Oct 2026 11:33:44 +0000 Subject: [PATCH 10/11] copilot: move readable model names to a child PR Co-authored-by: Isaac --- src/ucode/agents/claude.py | 29 ++++++++++++++++++++++++----- src/ucode/agents/copilot.py | 3 +-- src/ucode/databricks.py | 23 ----------------------- tests/test_agent_copilot.py | 4 ++-- 4 files changed, 27 insertions(+), 32 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 4b229a9f9..251d92fb7 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -48,7 +48,6 @@ build_tool_base_url, get_databricks_token, maybe_add_gpt_1m_suffix, - model_display_label, ug_binary, ) from ucode.launcher import exec_or_spawn @@ -603,7 +602,7 @@ def render_overlay( overlay["modelPicker"] = { "replaceBuiltInOptions": True, "options": [ - {"model": maybe_add_gpt_1m_suffix(m), "label": model_display_label(m)} + {"model": maybe_add_gpt_1m_suffix(m), "label": _picker_label(m)} for m in static_models ], } @@ -614,8 +613,7 @@ def render_overlay( "options": [ _picker_option( maybe_add_gpt_1m_suffix(model), - picker_catalog.model_id_to_display_name.get(model) - or model_display_label(model), + picker_catalog.model_id_to_display_name.get(model) or _picker_label(model), picker_catalog.model_id_to_description.get(model), ) for model in picker_catalog.model_ids @@ -634,11 +632,32 @@ def render_overlay( return overlay, keys +_MODEL_LABEL_ACRONYMS = frozenset({"glm", "gpt"}) _CLAUDE_PICKER_LABEL_RE = re.compile( r"^Claude (Fable|Opus|Sonnet|Haiku) (\d+(?:\.\d+)*)$", re.IGNORECASE ) +def _picker_label(model: str) -> str: + """A human-friendly picker label for a model id (e.g. ``system.ai.claude-haiku-4-5`` -> + ``Claude Haiku 4.5``): keep the vendor and name words title-cased, uppercase known acronyms, + and join a run of numeric segments into a dotted version.""" + stem = model.removeprefix("system.ai.") + parts: list[str] = [] + version: list[str] = [] + for token in stem.split("-"): + if token.isdigit(): + version.append(token) + continue + if version: + parts.append(".".join(version)) + version = [] + parts.append(token.upper() if token in _MODEL_LABEL_ACRONYMS else token.title()) + if version: + parts.append(".".join(version)) + return " ".join(parts) if parts else stem + + def _picker_option(model: str, label: str, description: str | None = None) -> dict[str, str]: option = {"model": model, "label": label} if description: @@ -681,7 +700,7 @@ def default_model_picker_catalog( descriptions: dict[str, str] = {} for family, raw_model in defaults.items(): model = raw_model - label = model_display_label(model.removesuffix("[1m]")) + label = _picker_label(model.removesuffix("[1m]")) if provider is not None: # Family shortcuts stay distinct from catalog rows for the same target. model = family diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 236505e4e..51138fcb0 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -49,7 +49,6 @@ ModelServiceEntry, build_copilot_base_url, build_tool_base_url, - model_display_label, model_version_sort_key, ) from ucode.os_compatibility import subprocess_cross_os @@ -449,7 +448,7 @@ def render_providers_config( "id": model.id, "provider": model.provider, "wireModel": model.id, - "name": model_display_label(model.id), + "name": model.id, } catalog_id = copilot_catalog_model_id(model.id) if model.wire == "anthropic" else None if catalog_id is not None: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 157f9110a..8c5321764 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3512,29 +3512,6 @@ def fetch_ai_gateway_claude_models(workspace: str, token: str) -> dict[str, str] return models -_MODEL_LABEL_ACRONYMS = frozenset({"glm", "gpt"}) - - -def model_display_label(model: str) -> str: - """A human-friendly picker label for a model id (e.g. ``system.ai.claude-haiku-4-5`` -> - ``Claude Haiku 4.5``): keep the vendor and name words title-cased, uppercase known acronyms, - and join a run of numeric segments into a dotted version.""" - stem = model.removeprefix("system.ai.") - parts: list[str] = [] - version: list[str] = [] - for token in stem.split("-"): - if token.isdigit(): - version.append(token) - continue - if version: - parts.append(".".join(version)) - version = [] - parts.append(token.upper() if token in _MODEL_LABEL_ACRONYMS else token.title()) - if version: - parts.append(".".join(version)) - return " ".join(parts) if parts else stem - - def model_version_sort_key(name: str) -> tuple: """Sort endpoint names so newer model versions come first. diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 248f5135c..c101fa121 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -342,12 +342,12 @@ def test_models_carry_claude_settings_and_limits(self): "id": SONNET, "provider": "dbx-anthropic", "wireModel": SONNET, - "name": "Claude Sonnet 5.5", + "name": SONNET, "modelId": "claude-sonnet-5.5", "maxPromptTokens": 872_000, "maxOutputTokens": 128_000, }, - {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": "GPT 5.4"}, + {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": GPT}, ] From 5ec1a9bc4da96e48de8c2587c1e5c4b3f93a0ea8 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 21 Oct 2026 11:33:46 +0000 Subject: [PATCH 11/11] copilot: readable model names in the session registry Co-authored-by: Isaac --- src/ucode/agents/claude.py | 29 +++++------------------------ src/ucode/agents/copilot.py | 3 ++- src/ucode/databricks.py | 23 +++++++++++++++++++++++ tests/test_agent_copilot.py | 4 ++-- 4 files changed, 32 insertions(+), 27 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 251d92fb7..4b229a9f9 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -48,6 +48,7 @@ build_tool_base_url, get_databricks_token, maybe_add_gpt_1m_suffix, + model_display_label, ug_binary, ) from ucode.launcher import exec_or_spawn @@ -602,7 +603,7 @@ def render_overlay( overlay["modelPicker"] = { "replaceBuiltInOptions": True, "options": [ - {"model": maybe_add_gpt_1m_suffix(m), "label": _picker_label(m)} + {"model": maybe_add_gpt_1m_suffix(m), "label": model_display_label(m)} for m in static_models ], } @@ -613,7 +614,8 @@ def render_overlay( "options": [ _picker_option( maybe_add_gpt_1m_suffix(model), - picker_catalog.model_id_to_display_name.get(model) or _picker_label(model), + picker_catalog.model_id_to_display_name.get(model) + or model_display_label(model), picker_catalog.model_id_to_description.get(model), ) for model in picker_catalog.model_ids @@ -632,32 +634,11 @@ def render_overlay( return overlay, keys -_MODEL_LABEL_ACRONYMS = frozenset({"glm", "gpt"}) _CLAUDE_PICKER_LABEL_RE = re.compile( r"^Claude (Fable|Opus|Sonnet|Haiku) (\d+(?:\.\d+)*)$", re.IGNORECASE ) -def _picker_label(model: str) -> str: - """A human-friendly picker label for a model id (e.g. ``system.ai.claude-haiku-4-5`` -> - ``Claude Haiku 4.5``): keep the vendor and name words title-cased, uppercase known acronyms, - and join a run of numeric segments into a dotted version.""" - stem = model.removeprefix("system.ai.") - parts: list[str] = [] - version: list[str] = [] - for token in stem.split("-"): - if token.isdigit(): - version.append(token) - continue - if version: - parts.append(".".join(version)) - version = [] - parts.append(token.upper() if token in _MODEL_LABEL_ACRONYMS else token.title()) - if version: - parts.append(".".join(version)) - return " ".join(parts) if parts else stem - - def _picker_option(model: str, label: str, description: str | None = None) -> dict[str, str]: option = {"model": model, "label": label} if description: @@ -700,7 +681,7 @@ def default_model_picker_catalog( descriptions: dict[str, str] = {} for family, raw_model in defaults.items(): model = raw_model - label = _picker_label(model.removesuffix("[1m]")) + label = model_display_label(model.removesuffix("[1m]")) if provider is not None: # Family shortcuts stay distinct from catalog rows for the same target. model = family diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 51138fcb0..236505e4e 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -49,6 +49,7 @@ ModelServiceEntry, build_copilot_base_url, build_tool_base_url, + model_display_label, model_version_sort_key, ) from ucode.os_compatibility import subprocess_cross_os @@ -448,7 +449,7 @@ def render_providers_config( "id": model.id, "provider": model.provider, "wireModel": model.id, - "name": model.id, + "name": model_display_label(model.id), } catalog_id = copilot_catalog_model_id(model.id) if model.wire == "anthropic" else None if catalog_id is not None: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 8c5321764..157f9110a 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3512,6 +3512,29 @@ def fetch_ai_gateway_claude_models(workspace: str, token: str) -> dict[str, str] return models +_MODEL_LABEL_ACRONYMS = frozenset({"glm", "gpt"}) + + +def model_display_label(model: str) -> str: + """A human-friendly picker label for a model id (e.g. ``system.ai.claude-haiku-4-5`` -> + ``Claude Haiku 4.5``): keep the vendor and name words title-cased, uppercase known acronyms, + and join a run of numeric segments into a dotted version.""" + stem = model.removeprefix("system.ai.") + parts: list[str] = [] + version: list[str] = [] + for token in stem.split("-"): + if token.isdigit(): + version.append(token) + continue + if version: + parts.append(".".join(version)) + version = [] + parts.append(token.upper() if token in _MODEL_LABEL_ACRONYMS else token.title()) + if version: + parts.append(".".join(version)) + return " ".join(parts) if parts else stem + + def model_version_sort_key(name: str) -> tuple: """Sort endpoint names so newer model versions come first. diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index c101fa121..248f5135c 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -342,12 +342,12 @@ def test_models_carry_claude_settings_and_limits(self): "id": SONNET, "provider": "dbx-anthropic", "wireModel": SONNET, - "name": SONNET, + "name": "Claude Sonnet 5.5", "modelId": "claude-sonnet-5.5", "maxPromptTokens": 872_000, "maxOutputTokens": 128_000, }, - {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": GPT}, + {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": "GPT 5.4"}, ]