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/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 40c33f386..d08f7eba5 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,25 @@ write_dotenv, write_json_file, ) +from ucode.constants import AI_GATEWAY_REQUEST_TAGS_HEADER, LOOPBACK_HOST +from ucode.custom_oauth import launch_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 +77,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 +146,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 +229,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 +263,85 @@ 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 _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 _token_provider(state: dict) -> Callable[[bool], str]: + workspace = state["workspace"] + return lambda force_refresh: launch_token(state, workspace, force_refresh=force_refresh) -def _refresh_forever(state: dict, model: str, stop_event: threading.Event) -> None: - while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS): +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 _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 _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() + server.server_close() + 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/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/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..ba51b2caa 100644 --- a/tests/README.md +++ b/tests/README.md @@ -432,6 +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_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) | | `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..07e46b35c 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -203,6 +203,15 @@ 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. `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, 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 `utils/agents/claude.py` and `utils/agents/codex.py` own the existing native answer @@ -256,6 +265,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 +424,8 @@ 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_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 @@ -607,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. 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/` @@ -833,6 +845,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`, 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 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..cae6f08a7 --- /dev/null +++ b/tests/integration/test_ug_copilot_headless.py @@ -0,0 +1,77 @@ +"""CUJs for using GitHub Copilot CLI from scripts through installed ug.""" + +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 + +pytestmark = [pytest.mark.live, pytest.mark.copilot] + + +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_bearer_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_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..42bd005bc --- /dev/null +++ b/tests/integration/test_ug_copilot_lifecycle.py @@ -0,0 +1,265 @@ +"""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" +# 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): + 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 _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()) + + +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) + 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)} + ) + 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. + """ + 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() + + +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) 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 a7899bfe2..0e6e32c55 100644 --- a/tests/integration/utils/evidence.py +++ b/tests/integration/utils/evidence.py @@ -132,6 +132,21 @@ def assistant_answer_contains(session, agent: str, value: str, *, child: bool = ) +def assert_no_bearer_on_disk(session) -> None: + """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." + 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") + 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] @@ -215,6 +230,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/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() 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 430b3563e..71e995ec4 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -3,10 +3,11 @@ from __future__ import annotations import json -import threading +import httpx import pytest +from ucode import custom_oauth from ucode.agents import copilot WS = "https://example.databricks.com" @@ -28,11 +29,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 +42,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 +79,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 +126,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 +348,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 +365,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(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 + + 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(custom_oauth, "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(custom_oauth, "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(custom_oauth, "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 +558,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 +569,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() - - def get_token(*args, force_refresh=False, **kwargs): - token_calls.append(force_refresh) - if force_refresh: - refreshed.set() - return "tok" - - monkeypatch.setattr(copilot, "get_databricks_token", get_token) - - class Process: - def wait(self): - assert refreshed.wait(timeout=2) - return 0 - - monkeypatch.setattr( - copilot.subprocess_cross_os, - "popen", - lambda *args, **kwargs: Process(), - ) + seen = _stub_launch(monkeypatch, token="dbx-token") + + with pytest.raises(SystemExit): + copilot.launch( + {"workspace": WS}, [], options=copilot.LaunchOptions(user_pinned_model="gpt-5") + ) + + 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 + + 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) + + 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(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"] + 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=10) + + 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 +650,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_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 adc406947..92d86a216 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,42 @@ 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() + server.server_close() + 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")