diff --git a/README.md b/README.md index 93b77b8ae..2781f8686 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,14 @@ restart, and no Databricks token is written to disk. If the Databricks sign-in i 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`. +`ug copilot --provider ..` routes the session through a Model Provider +Service instead, and `/model` lists that service's models: Claude models of an Anthropic or +Bedrock service, and the chat models of an OpenAI or Azure OpenAI service. For a service that +allows every model, `/model` lists the models the gateway reports for it, which the provider +account may not serve; name the one to start on with `--model`. Services that relay a Claude subscription work only with Claude Code. Interactive +`ug configure` offers the matching services under **External Models**, and a managed workspace +config can set one for Copilot. + 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/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index e507187fa..af8a2d5c5 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -468,6 +468,11 @@ def configure_tool( parent_schema=parent_schema, picker_catalog=picker_catalog, ) + elif tool == "copilot": + # Under a model provider service, Copilot starts on one of the service's models, chosen at launch. + if not model and not provider: + raise RuntimeError(f"A {tool} model must be selected before configuration.") + result = copilot.write_tool_config(state, model) else: # Every tool in this branch needs a model — including gemini under a provider, # which still pins the service's target model in the URL. @@ -475,8 +480,6 @@ def configure_tool( raise RuntimeError(f"A {tool} model must be selected before configuration.") if tool == "gemini": result = gemini.write_tool_config(state, model, provider=provider) - elif tool == "copilot": - result = copilot.write_tool_config(state, model) elif tool == "pi": result = pi.write_tool_config(state, model) elif tool == "kilo": diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 820a1b4a1..dcb1df2ef 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -10,6 +10,10 @@ models are excluded: Databricks' Gemini translation rejects the `stream_options` field Copilot sends, and open-source models are not validated with Copilot yet. +`ug copilot --provider` registers a Model Provider Service's models instead: its Claude models on +the Anthropic endpoint, its other chat models on the gateway's native OpenAI route, with every +registry provider sending the service's routing header. + Copilot holds its provider credential for the whole session, so every provider points at a per-session loopback proxy (`gateway_proxy.start_gateway_auth_proxy`): Copilot authenticates to the proxy with a random session secret, written only into the session registry, and the proxy @@ -42,14 +46,28 @@ write_dotenv, write_json_file, ) -from ucode.constants import AI_GATEWAY_REQUEST_TAGS_HEADER, LOOPBACK_HOST +from ucode.constants import ( + AI_GATEWAY_REQUEST_TAGS_HEADER, + API_ANTHROPIC_MESSAGES, + API_CODEX_RESPONSES, + API_MLFLOW_CHAT_COMPLETIONS, + API_MLFLOW_RESPONSES, + API_OPENAI_CHAT_COMPLETIONS, + API_OPENAI_RESPONSES, + LOOPBACK_HOST, + MODEL_PROVIDER_SERVICE_HEADER, +) from ucode.custom_oauth import launch_token from ucode.databricks import ( + ANTHROPIC_API_PROVIDER_TYPES, ANTHROPIC_FAMILIES, ModelServiceEntry, + bare_model_name, build_copilot_base_url, build_tool_base_url, + list_provider_service_models, model_version_sort_key, + resolve_provider_service, ) from ucode.os_compatibility import subprocess_cross_os from ucode.os_compatibility.file_lock_cross_os import ( @@ -60,7 +78,7 @@ from ucode.telemetry import agent_version, ucode_user_agent, ug_version from ucode.ui import print_warning -from .args import LaunchOptions +from .args import LaunchOptions, explicit_model_arg_value from .request_tags import request_tags_header_value COPILOT_CONFIG_DIR = Path.home() / ".copilot" @@ -96,16 +114,23 @@ # Earlier versions wrote the Databricks token (or PAT) under these keys. CREDENTIAL_ENV_KEYS = ("COPILOT_PROVIDER_BEARER_TOKEN", "OAUTH_TOKEN") CATALOG_STATE_KEY = "copilot_catalog" +# The Model Provider Service this launch routes through; set by the launch command, never saved. +LAUNCH_PROVIDER_SERVICE_KEY = "_copilot_launch_provider_service" + +# The wire APIs Copilot speaks; `responses` and `completions` are Copilot's `wireApi` values. +WIRE_ANTHROPIC = "anthropic" +WIRE_RESPONSES = "responses" +WIRE_COMPLETIONS = "completions" PROVIDER_NAMES: dict[str, str] = { - "anthropic": "dbx-anthropic", - "responses": "dbx-responses", - "completions": "dbx-completions", + WIRE_ANTHROPIC: "dbx-anthropic", + WIRE_RESPONSES: "dbx-responses", + WIRE_COMPLETIONS: "dbx-completions", } _WIRE_ROUTES: dict[str, tuple[str, ...]] = { - "anthropic": ("anthropic/v1/messages",), - "responses": ("mlflow/v1/responses", "openai/v1/responses", "codex/v1/responses"), - "completions": ("mlflow/v1/chat/completions", "openai/v1/chat/completions"), + WIRE_ANTHROPIC: (API_ANTHROPIC_MESSAGES,), + WIRE_RESPONSES: (API_MLFLOW_RESPONSES, API_OPENAI_RESPONSES, API_CODEX_RESPONSES), + WIRE_COMPLETIONS: (API_MLFLOW_CHAT_COMPLETIONS, API_OPENAI_CHAT_COMPLETIONS), } _UNSUPPORTED_MODELS = { "claude-opus-4-1": "the gateway rejects the thinking settings Copilot sends for it", @@ -113,7 +138,14 @@ _GATEWAY_ALIAS_PREFIX = "anthropic-aigw-" _GPT_MODEL_MAJOR_PATTERN = re.compile(r"^gpt-(\d+)o?(?=$|[.-])") _CLAUDE_DASHED_VERSION_PATTERN = re.compile(r"-(\d+)-(\d+)$") +# Anthropic and Bedrock Claude ids end in a release date and, on Bedrock, a version: +# `-20250929-v1:0`. +_CLAUDE_RELEASE_SUFFIX_PATTERN = re.compile(r"(?:-\d{8})?(?:-v\d+(?::\d+)?)?$") _CLAUDE_HAIKU_MAJOR_PATTERN = re.compile(r"^claude-haiku-(\d+)") +# Model Provider Service ids write versions with dots (`gpt-5.6`) or no separator (`o3`); the +# version sort reads dash-separated digits. +_DOTTED_VERSION_PATTERN = re.compile(r"(?<=\d)\.(?=\d)") +_O_SERIES_PATTERN = re.compile(r"^o(?=\d)") _CLAUDE_CATALOG_VERSION_PATTERN = re.compile(r"^claude-([a-z]+)-(\d+)(?:\.(\d+))?$") # (maxPromptTokens, maxOutputTokens) by Copilot catalog id. Copilot assumes a 128K prompt for @@ -135,57 +167,77 @@ } _CLAUDE_1M_LIMITS = (936_000, 32_000) _CLAUDE_200K_LIMITS = (168_000, 32_000) +_CLAUDE_200K_WINDOW = 200_000 - -def _model_tail(model: str) -> str: - """The model name without its ``..`` and ``databricks-`` prefixes or ``[...]`` suffix.""" - name = model.split("[", 1)[0].lower() - tail = name.split(".", 2)[-1] if name.count(".") >= 2 else name - return tail.removeprefix("databricks-") +# A Model Provider Service target that declares its APIs is usable only if one is a chat API. +_CHAT_API_TYPES = frozenset(route for routes in _WIRE_ROUTES.values() for route in routes) def copilot_exclusion_reason(model: str, api_types: Iterable[str] = ()) -> str | None: """Why ``model`` can't be offered in Copilot, or None when it can.""" - return _classify(model, tuple(api_types))[1] + return classify_model(model, api_types)[1] def classify_copilot_model(model: str, api_types: Iterable[str] = ()) -> str | None: - """The wire API (``anthropic``, ``responses`` or ``completions``) Copilot uses for ``model``. - - Names decide the wire; ``api_types`` (the model service's advertised routes) only excludes, - because the gateway serves some models on routes they don't list.""" - return _classify(model, tuple(api_types))[0] - - -def _classify(model: str, api_types: tuple[str, ...]) -> tuple[str | None, str | None]: - if model.lower().startswith(_GATEWAY_ALIAS_PREFIX): - return None, "it is a gateway alias, not a model" - tail = _model_tail(model) - if tail in _UNSUPPORTED_MODELS: - return None, _UNSUPPORTED_MODELS[tail] - gpt = _GPT_MODEL_MAJOR_PATTERN.match(tail) - if tail.startswith("claude-"): - wire = "anthropic" - elif gpt: - wire = "responses" if int(gpt.group(1)) >= 5 else "completions" - elif "gemini-" in tail: - return None, "the gateway rejects the requests Copilot sends to Gemini models" - else: + """The wire API (``anthropic``, ``responses`` or ``completions``) Copilot uses for ``model``.""" + return classify_model(model, api_types)[0] + + +def classify_model( + model: str, api_types: Iterable[str] = (), *, provider_type: str | None = None +) -> tuple[str | None, str | None]: + """The wire API Copilot uses for ``model``, or why it can't use it. + + Names decide the wire, so ``api_types`` only excludes: the gateway serves some models on APIs + they don't declare. Without ``provider_type``, ``model`` is a Databricks model service and + ``api_types`` its advertised routes. With it, ``model`` is a Model Provider Service target and + ``api_types`` its declared APIs: Anthropic-API provider types serve only Claude, and other types + put every model that isn't GPT 5+ or o-series on Chat Completions.""" + declared = tuple(api_types) + tail = bare_model_name(model) + if provider_type is None: + if model.lower().startswith(_GATEWAY_ALIAS_PREFIX): + return None, "it is a gateway alias, not a model" + if tail in _UNSUPPORTED_MODELS: + return None, _UNSUPPORTED_MODELS[tail] + elif declared and _CHAT_API_TYPES.isdisjoint(declared): + return None, f"it isn't a chat model (it serves {', '.join(declared)})" + wire = _wire_from_name(tail) + if provider_type in ANTHROPIC_API_PROVIDER_TYPES: + if wire != WIRE_ANTHROPIC: + return None, "ug copilot uses only the Claude models of this service" + return wire, None + if provider_type is not None: + return (WIRE_RESPONSES if wire == WIRE_RESPONSES else WIRE_COMPLETIONS), None + if wire is None: + if "gemini-" in tail: + return None, "the gateway rejects the requests Copilot sends to Gemini models" return None, "ug copilot supports Claude and GPT models" - if api_types and not set(api_types) & set(_WIRE_ROUTES[wire]): + if declared and not set(declared) & set(_WIRE_ROUTES[wire]): return None, f"it isn't served on {_WIRE_ROUTES[wire][0]}" return wire, None +def _wire_from_name(name: str) -> str | None: + if name.startswith("claude-"): + return WIRE_ANTHROPIC + gpt = _GPT_MODEL_MAJOR_PATTERN.match(name) + # Copilot sends `temperature: 0` on Chat Completions, which GPT 5+ and o-series models reject. + if (gpt and int(gpt.group(1)) >= 5) or _O_SERIES_PATTERN.match(name): + return WIRE_RESPONSES + return WIRE_COMPLETIONS if gpt else None + + def copilot_catalog_model_id(model: str) -> str | None: """Copilot's catalog id for a Claude model, e.g. ``system.ai.claude-sonnet-5-5`` -> ``claude-sonnet-5.5``. Copilot applies Claude request settings to any ``claude-*`` id, so ids missing from its catalog still work. Returns None for anything other than Claude. """ - tail = _model_tail(model) + tail = bare_model_name(model) if not tail.startswith("claude-"): return None + tail = _CLAUDE_RELEASE_SUFFIX_PATTERN.sub("", tail, count=1) catalog_id = _CLAUDE_DASHED_VERSION_PATTERN.sub(r"-\1.\2", tail) # Copilot sends budget thinking for any Haiku it doesn't know, which Haiku 5+ rejects; Sonnet # ids get the adaptive thinking those models require. @@ -232,22 +284,93 @@ def _catalog_sort_key(model: CopilotModel) -> tuple: ( i for i, f in enumerate(ANTHROPIC_FAMILIES) - if _model_tail(model.id).startswith(f"claude-{f}-") + if bare_model_name(model.id).startswith(f"claude-{f}-") ), len(ANTHROPIC_FAMILIES), ) - return (tuple(PROVIDER_NAMES).index(model.wire), family, model_version_sort_key(model.id)) + sort_name = _O_SERIES_PATTERN.sub("o-", _DOTTED_VERSION_PATTERN.sub("-", model.id)) + version_key = model_version_sort_key(sort_name) + return (tuple(PROVIDER_NAMES).index(model.wire), family, version_key) + + +def build_catalog( + models: Iterable[tuple[str, Iterable[str]]], *, provider_type: str | None = None +) -> tuple[list[CopilotModel], list[tuple[str, str]]]: + """The ``(model, api_types)`` pairs Copilot can use, sorted Claude first (by family, newest + first) then GPT, and the skipped models with the reason (see :func:`classify_model`).""" + usable: dict[str, CopilotModel] = {} + skipped: dict[str, str] = {} + for model, api_types in models: + wire, reason = classify_model(model, api_types, provider_type=provider_type) + if wire is not None: + usable.setdefault(model, CopilotModel(model, wire)) + elif reason is not None: + skipped.setdefault(model, reason) + return sorted(usable.values(), key=_catalog_sort_key), list(skipped.items()) def build_copilot_catalog(entries: Iterable[ModelServiceEntry | str]) -> list[CopilotModel]: - """The models Copilot can use, Claude first (by family, newest first), then GPT.""" - models: dict[str, CopilotModel] = {} - for entry in entries: - model, api_types = (entry, ()) if isinstance(entry, str) else (entry.id, entry.api_types) - wire = classify_copilot_model(model, api_types) - if wire is not None and model not in models: - models[model] = CopilotModel(model, wire) - return sorted(models.values(), key=_catalog_sort_key) + """The Databricks models Copilot can use (see :func:`build_catalog`).""" + pairs = ((e, ()) if isinstance(e, str) else (e.id, e.api_types) for e in entries) + return build_catalog(pairs)[0] + + +def _without_registry_prefix(model: str) -> str: + """``model`` without a Copilot registry provider prefix: ``dbx-responses/gpt-5.5`` -> ``gpt-5.5``.""" + name, separator, rest = model.partition("/") + return rest if separator and name in PROVIDER_NAMES.values() else model + + +def provider_service_catalog( + state: dict, service_name: str, requested: str | None +) -> list[CopilotModel]: + """The models of Model Provider Service ``service_name`` that this launch registers. + + A service that allows every model registers what the gateway lists for it, plus ``requested`` + (the launch's ``--model``) or the managed default, so a model the listing omits can be named.""" + workspace = state["workspace"] + token = launch_token(state, workspace) + service, error = resolve_provider_service("copilot", service_name, workspace, token) + if error or service is None: + raise RuntimeError(error or f"Model provider service '{service_name}' was not found.") + provider_type = service["provider_type"] + named = list(service.get("targets") or []) + listed: list[str] = [] + allow_all = bool(service.get("allow_all_targets")) + if allow_all: + start = _without_registry_prefix(requested) if requested else None + managed_default = state.get("copilot_default_model") + start = start or (managed_default if isinstance(managed_default, str) else None) + listed, reason = list_provider_service_models(workspace, token, service_name, provider_type) + if reason and not start: + raise RuntimeError( + f"Could not list the models of '{service_name}': {reason}. Name the one to start " + "on with `ug copilot --model `." + ) + if reason: + print_warning(f"Could not list the models of {service_name}: {reason}.") + if start: + named.append(start) + api_types = service.get("target_api_types") or {} + catalog, skipped = build_catalog( + ((model, api_types.get(model, ())) for model in [*named, *listed]), + provider_type=provider_type, + ) + # An allow-all listing includes models Copilot can't use; report only declared or named ones. + skipped = [(model, reason) for model, reason in skipped if model in named] + if not catalog: + detail = "".join(f" {model} ({reason})." for model, reason in skipped) + raise RuntimeError( + f"Model provider service '{service_name}' has no models ug copilot can use.{detail}" + + (" Name one it serves with `ug copilot --model `." if allow_all else "") + ) + if skipped: + print_warning( + f"Skipping models of {service_name} that ug copilot can't use: " + + "; ".join(f"{model} ({reason})" for model, reason in skipped) + + "." + ) + return catalog def catalog_from_listings( @@ -264,8 +387,10 @@ def catalog_from_listings( return fallback_catalog or None catalog = build_copilot_catalog(listed) # A family the UC listing lacks keeps the models its gateway listing found. - families = {model.wire == "anthropic" for model in catalog} - missing = [model for model in fallback_catalog if (model.wire == "anthropic") not in families] + families = {model.wire == WIRE_ANTHROPIC for model in catalog} + missing = [ + model for model in fallback_catalog if (model.wire == WIRE_ANTHROPIC) not in families + ] return sorted([*catalog, *missing], key=_catalog_sort_key) @@ -362,28 +487,33 @@ def _first_available( def launch_model( - state: dict, catalog: list[CopilotModel], *, requested: str | None = None, warn: bool = True + state: dict, + catalog: list[CopilotModel], + *, + requested: str | None = None, + warn: bool = True, + managed_list: bool = True, ) -> CopilotModel: """The model a session starts on: ``requested`` (budget recommendation or managed default), - then the managed default and list, then the newest Sonnet, Opus or Haiku, then GPT.""" - servable_static = ( - m for m in _string_list(state.get("copilot_static_models")) if classify_copilot_model(m) - ) - managed_models = _string_list(state.get("copilot_models")) + then the managed default and list, then the newest Sonnet, Opus or Haiku, then GPT. + + ``managed_list=False`` ignores the managed model lists, which name Databricks models.""" + static = _string_list(state.get("copilot_static_models")) if managed_list else [] + managed_models = _string_list(state.get("copilot_models")) if managed_list else [] candidates = [ requested, state.get("copilot_default_model"), - next(servable_static, None), + next((m for m in static if classify_copilot_model(m)), None), managed_models[0] if managed_models else None, ] chosen = _first_available(dict.fromkeys(c for c in candidates if c), catalog, warn=warn) if chosen is not None: return chosen for prefix in ("claude-sonnet-", "claude-opus-", "claude-haiku-"): - family = next((m for m in catalog if _model_tail(m.id).startswith(prefix)), None) + family = next((m for m in catalog if bare_model_name(m.id).startswith(prefix)), None) if family is not None: return family - return next((m for m in catalog if m.wire != "anthropic"), catalog[0]) + return next((m for m in catalog if m.wire != WIRE_ANTHROPIC), catalog[0]) def default_model(state: dict) -> str | None: @@ -426,17 +556,24 @@ def rewrite_model_args( def render_providers_config( - catalog: list[CopilotModel], proxy_base_url: str, session_secret: str + catalog: list[CopilotModel], + proxy_base_url: str, + session_secret: str, + *, + provider_service: str | None = None, ) -> dict: """Copilot's providers.json for ``catalog``: one provider per wire, all on the session proxy. A provider with no models makes Copilot probe ``GET /models``, so only used wires are - listed.""" + listed. With ``provider_service`` (a Model Provider Service), every registry provider carries its + routing header and OpenAI wires use the gateway's native OpenAI route, which serves Model + Provider Services untranslated.""" providers = [] for wire, name in PROVIDER_NAMES.items(): if not any(model.wire == wire for model in catalog): continue - if wire == "anthropic": + provider: dict[str, Any] + if wire == WIRE_ANTHROPIC: provider = { "name": name, "type": "anthropic", @@ -447,10 +584,16 @@ def render_providers_config( "name": name, "type": "openai", "wireApi": wire, - "baseUrl": build_copilot_base_url(proxy_base_url), + "baseUrl": ( + f"{proxy_base_url}/ai-gateway/openai/v1" + if provider_service + else build_copilot_base_url(proxy_base_url) + ), } # `bearerToken` reaches the proxy as `Authorization` on every provider type. provider["bearerToken"] = session_secret + if provider_service: + provider["headers"] = {MODEL_PROVIDER_SERVICE_HEADER: provider_service} providers.append(provider) models = [] for model in catalog: @@ -460,10 +603,13 @@ def render_providers_config( "wireModel": model.id, "name": model.id, } - catalog_id = copilot_catalog_model_id(model.id) if model.wire == "anthropic" else None + catalog_id = copilot_catalog_model_id(model.id) if model.wire == WIRE_ANTHROPIC else None if catalog_id is not None: entry["modelId"] = catalog_id - limits = copilot_model_limits(model.id) if model.wire == "anthropic" else None + limits = copilot_model_limits(model.id) if model.wire == WIRE_ANTHROPIC else None + if limits is not None and provider_service and sum(limits) > _CLAUDE_200K_WINDOW: + # A Model Provider Service may not serve the 1M window. + limits = _CLAUDE_200K_LIMITS if limits is not None: entry["maxPromptTokens"], entry["maxOutputTokens"] = limits models.append(entry) @@ -630,17 +776,20 @@ def scrub_stored_credentials() -> None: write_dotenv(COPILOT_ENV_PATH, existing) -def write_tool_config(state: dict, model: str) -> dict: +def write_tool_config(state: dict, model: str | None) -> 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) for key in LEGACY_ENV_KEYS: existing.pop(key, None) - try: - existing["COPILOT_MODEL"] = resolve_copilot_model(model, copilot_catalog(state)).ref - except RuntimeError: - existing["COPILOT_MODEL"] = model + if model is None: + existing.pop("COPILOT_MODEL", None) + else: + try: + existing["COPILOT_MODEL"] = resolve_copilot_model(model, copilot_catalog(state)).ref + except RuntimeError: + existing["COPILOT_MODEL"] = model existing["COPILOT_OFFLINE"] = "true" write_dotenv(COPILOT_ENV_PATH, existing) state = mark_tool_managed(state, "copilot", MANAGED_KEYS) @@ -696,18 +845,34 @@ def _wait_for_exit(proc: _ChildProcess) -> int: continue -def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: +def _session_catalog( + state: dict, tool_args: list[str], provider_service: str | None +) -> list[CopilotModel]: + """The models this launch registers: the Model Provider Service's, else the Databricks + catalog.""" + if provider_service: + return provider_service_catalog( + state, provider_service, explicit_model_arg_value(tool_args) + ) catalog = copilot_catalog(state) if not catalog: raise RuntimeError( - "No Copilot model is available on this workspace: ug copilot supports its Claude and " - "GPT models. Run `ug configure --agents copilot` to rediscover them." + "No Copilot model is available on this workspace: ug copilot supports its Claude " + "and GPT models. Run `ug configure --agents copilot` to rediscover them." ) for model, reason in unservable_static_models(state): print_warning(f"Skipping {model} from your workspace's managed config: {reason}.") + return catalog + + +def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: + provider_service = state.pop(LAUNCH_PROVIDER_SERVICE_KEY, None) + catalog = _session_catalog(state, tool_args, provider_service) # Resolve before the proxy starts so a bad model fails without starting Copilot. tool_args, selected = rewrite_model_args(tool_args, catalog) - start_model = selected or launch_model(state, catalog, requested=options.launch_model) + start_model = selected or launch_model( + state, catalog, requested=options.launch_model, managed_list=not provider_service + ) scrub_stored_credentials() # Only Copilot launches need the proxy (and its `httpx` import). @@ -722,11 +887,16 @@ def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None 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, + # Copilot sends `reasoning_effort` on Chat Completions, but only models that aren't GPT 5+ + # or o-series use that wire, and OpenAI and Azure reject the field for them. + dropped_chat_completions_fields=("reasoning_effort",), ) threading.Thread(target=server.serve_forever, daemon=True).start() try: proxy_base_url = f"http://{LOOPBACK_HOST}:{server.server_address[1]}" - config = render_providers_config(catalog, proxy_base_url, session_secret) + config = render_providers_config( + catalog, proxy_base_url, session_secret, provider_service=provider_service + ) wire_override = os.environ.get("COPILOT_PROVIDER_WIRE_MODEL") if wire_override: for entry in config["models"]: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 49ee98ae7..b084e1e0d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -889,15 +889,15 @@ def _provider_summary(tool: str, state: dict) -> str: def _maybe_select_provider_service(tool: str, state: dict) -> dict: - """Interactively let the user pick where claude/codex get their models — Databricks hosted + """Interactively let the user pick where an agent gets its models — Databricks hosted (``system.ai``), a Model Provider Service, or a custom Unity Catalog model location — and persist (or clear) the choice. - No-op for tools other than claude/codex/gemini (gemini has no custom location). Falls back - to Databricks when no matching provider services are found or the listing fails and no - custom location is offered. + No-op for tools other than claude/codex/gemini/copilot (gemini and copilot have no custom + location). Falls back to Databricks when no matching provider services are found or the + listing fails and no custom location is offered. """ - if tool not in ("claude", "codex", "gemini"): + if tool not in ("claude", "codex", "gemini", "copilot"): return state display = TOOL_SPECS[tool]["display"] @@ -3654,6 +3654,8 @@ def _launch_tool( state["_codex_launch_provider"] = provider elif parent_schema: state["_codex_launch_parent_schema"] = parent_schema + elif tool == "copilot" and provider: + state[copilot_agent.LAUNCH_PROVIDER_SERVICE_KEY] = provider launch_options = _launch_options( tool, ctx.args, @@ -4104,10 +4106,25 @@ def kilo_cmd( ) def copilot_cmd( ctx: typer.Context, + provider: Annotated[ + str | None, + typer.Option( + "--provider", + help="Route through a Unity Catalog Model Provider Service " + "(..); `/model` then lists that service's models.", + ), + ] = None, skip_preflight: SkipPreflightOption = False, + workspace: WorkspaceOption = None, ) -> None: """Launch GitHub Copilot CLI via Databricks.""" - _launch_tool("copilot", ctx, skip_preflight=skip_preflight) + _launch_tool( + "copilot", + ctx, + provider=provider, + skip_preflight=skip_preflight, + workspace_url=workspace, + ) @app.command( diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 7c2efb21d..3a5d7c385 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -20,6 +20,16 @@ ) MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service" + +# AI Gateway API types, as a model service lists them in `supported_api_types` and a Model +# Provider Service target in `native_api_types`. +API_ANTHROPIC_MESSAGES = "anthropic/v1/messages" +API_CODEX_RESPONSES = "codex/v1/responses" +API_GEMINI_GENERATE_CONTENT = "gemini/v1/generateContent" +API_MLFLOW_CHAT_COMPLETIONS = "mlflow/v1/chat/completions" +API_MLFLOW_RESPONSES = "mlflow/v1/responses" +API_OPENAI_CHAT_COMPLETIONS = "openai/v1/chat/completions" +API_OPENAI_RESPONSES = "openai/v1/responses" MODEL_SERVICE_PARENT_SCHEMA_HEADER = "Databricks-Model-Service-Parent-Schema" # Names the smart-router recipe (e.g. `task_v3`) in use; sent only when smart routing is enabled. SMART_ROUTER_RECIPE_HEADER = "Databricks-Smart-Router-Recipe" diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d3a50f83d..f916243d9 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -32,8 +32,11 @@ from urllib import request as urllib_request from urllib.parse import quote, urlencode, urlparse +from ucode.codex_config import catalog_slugs from ucode.config_io import APP_DIR from ucode.constants import ( + API_GEMINI_GENERATE_CONTENT, + API_OPENAI_RESPONSES, MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, ) @@ -114,6 +117,10 @@ class CodexMpsModelCatalogUnavailable(RuntimeError): """The workspace does not expose the Codex model-catalog route.""" +class CodexModelCatalogEmpty(RuntimeError): + """The Codex model-catalog route answered with no models.""" + + class CodexCatalogSource(Enum): PROVIDER = (MODEL_PROVIDER_SERVICE_HEADER, "Provider") PARENT_SCHEMA = (MODEL_SERVICE_PARENT_SCHEMA_HEADER, "Parent schema") @@ -2756,19 +2763,27 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: # the `openai`, `azure_openai`, `microsoft_foundry`, `amazon_bedrock`, and # `bedrock_mantle` provider types all serve (the gateway fronts them with the # OpenAI surface); gemini speaks Google's, served by a Gemini Enterprise -# provider. Tags are the short form produced by `_provider_type_tag` (e.g. -# `amazon_bedrock`). +# provider. copilot speaks Anthropic's for the Claude targets of anthropic and +# Bedrock services, and OpenAI's for openai and azure_openai services, on the +# gateway's native route for each API. Tags are the short form produced by `_provider_type_tag` +# (e.g. `amazon_bedrock`). _TOOL_PROVIDER_TYPES: dict[str, tuple[str, ...]] = { "claude": ("anthropic", "amazon_bedrock", "bedrock_mantle"), "codex": ("openai", "azure_openai", "microsoft_foundry", "amazon_bedrock", "bedrock_mantle"), "gemini": ("gemini_enterprise",), + "copilot": ("anthropic", "amazon_bedrock", "bedrock_mantle", "openai", "azure_openai"), } +# Agents that can only use a Bedrock service's Claude targets. +_CLAUDE_ONLY_BEDROCK_TOOLS = ("claude", "copilot") + # Provider types that expose Bedrock-style model ids (e.g. # `us.anthropic.claude-sonnet-4-6`, or region-less `anthropic.claude-*` on # Mantle) instead of the agent's canonical model names, so ucode must pin them # explicitly. BEDROCK_PROVIDER_TYPES: tuple[str, ...] = ("amazon_bedrock", "bedrock_mantle") +# Provider types the gateway serves on its Anthropic Messages route. +ANTHROPIC_API_PROVIDER_TYPES: tuple[str, ...] = ("anthropic", *BEDROCK_PROVIDER_TYPES) def tool_supports_provider_type(tool: str, provider_type: str) -> bool: @@ -2798,9 +2813,11 @@ def list_model_provider_services( Returns ``(services, reason)`` where each service is ``{"name": "..", "provider_type": "anthropic"|..., - "targets": [model_id, ...], "allow_all_targets": bool, "relayed": bool}``. + "targets": [model_id, ...], "target_api_types": {model_id: (api_type, ...)}, + "allow_all_targets": bool, "relayed": bool}``. ``targets`` is the provider-side model ids the service exposes (used to pin - Bedrock model names). ``relayed`` is True for a credential-less Anthropic + Bedrock model names); ``target_api_types`` holds the APIs each target declares + (``native_api_types``, e.g. ``openai/v1/embeddings``). ``relayed`` is True for a credential-less Anthropic service (Claude Max/Team/Enterprise subscription relay). A non-None ``reason`` means the listing call itself failed. @@ -2880,13 +2897,18 @@ def _provider_service_entry(raw_service: object) -> dict | None: raw_config = service.get("config") config = cast("dict[str, object]", raw_config) if isinstance(raw_config, dict) else {} targets: list[str] = [] + target_api_types: dict[str, tuple[str, ...]] = {} raw_targets = config.get("targets") for target in raw_targets if isinstance(raw_targets, list) else []: if not isinstance(target, dict): continue - model_id = cast("dict[str, object]", target).get("model") + entry = cast("dict[str, object]", target) + model_id = entry.get("model") if isinstance(model_id, str) and model_id: targets.append(model_id) + api_types = entry.get("native_api_types") + if isinstance(api_types, list): + target_api_types[model_id] = tuple(t for t in api_types if isinstance(t, str)) # Relayed = credential-less Anthropic (subscription relay). Only whether # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed # server-side, so both launch identically. @@ -2897,6 +2919,7 @@ def _provider_service_entry(raw_service: object) -> dict | None: "name": full_name, "provider_type": _provider_type_tag(raw_type if isinstance(raw_type, str) else None), "targets": targets, + "target_api_types": target_api_types, "allow_all_targets": bool(config.get("allow_all_targets")), "relayed": relayed, } @@ -2982,23 +3005,44 @@ def list_tool_provider_services( return names, None -def service_usable_for_tool(tool: str, service: dict) -> bool: - """True when ``tool`` can actually route through ``service``. +def _provider_service_unusable_reason(tool: str, service: dict, service_name: str) -> str | None: + """Why ``tool`` can't route through ``service``, or None when it can. - Beyond the provider-type match, a Bedrock service is only usable for claude - if it exposes at least one Claude model in its targets — otherwise there's no - routable model id to pin — or is ``allow_all_targets``, in which case the model - id comes from elsewhere (e.g. the managed config's authored default). (Anthropic - services use canonical names, so any match is usable.) + Beyond the provider-type match, only claude can use a relayed (subscription) service, and + claude and copilot can use a Bedrock service only if it has a Claude target — otherwise there's + no routable model id to pin — or is ``allow_all_targets``, in which case the model id comes + from elsewhere (the managed config's default or ``--model``). (Anthropic services use canonical + names, so any match is usable.) """ provider_type = service.get("provider_type", "") if not tool_supports_provider_type(tool, provider_type): - return False - if tool == "claude" and provider_type in BEDROCK_PROVIDER_TYPES: - return bool(service.get("allow_all_targets")) or bool( - map_claude_family_models(service.get("targets") or []) + supported = ", ".join(_TOOL_PROVIDER_TYPES.get(tool, ())) or "none" + return ( + f"Model provider service '{service_name}' is a '{provider_type}' provider, " + f"which {tool} can't route to (supported: {supported})." ) - return True + if tool != "claude" and service.get("relayed"): + return ( + f"Model provider service '{service_name}' relays a Claude subscription, which only " + f"Claude Code can use. Pick a Model Provider Service with an API key for {tool}." + ) + if ( + tool in _CLAUDE_ONLY_BEDROCK_TOOLS + and provider_type in BEDROCK_PROVIDER_TYPES + and not service.get("allow_all_targets") + and not map_claude_family_models(service.get("targets") or []) + ): + return ( + f"Model provider service '{service_name}' exposes no Claude models — " + f"add Claude targets to it, enable allow_all_targets, or pick a different service." + ) + return None + + +def service_usable_for_tool(tool: str, service: dict) -> bool: + """True when ``tool`` can actually route through ``service`` (see + :func:`_provider_service_unusable_reason`).""" + return _provider_service_unusable_reason(tool, service, service.get("name", "")) is None def resolve_provider_service( @@ -3030,23 +3074,9 @@ def resolve_provider_service( suffix = f" Available for {tool}: {', '.join(usable)}." if usable else "" detail = f" ({get_reason})" if get_reason and "404" not in get_reason else "" return None, f"Model provider service '{service_name}' was not found.{detail}{suffix}" - provider_type = match["provider_type"] - if not tool_supports_provider_type(tool, provider_type): - supported = ", ".join(_TOOL_PROVIDER_TYPES.get(tool, ())) or "none" - return None, ( - f"Model provider service '{service_name}' is a '{provider_type}' provider, " - f"which {tool} can't route to (supported: {supported})." - ) - if ( - tool == "claude" - and provider_type in BEDROCK_PROVIDER_TYPES - and not match.get("allow_all_targets") - and not map_claude_family_models(match.get("targets") or []) - ): - return None, ( - f"Model provider service '{service_name}' exposes no Claude models — " - f"add Claude targets to it, enable allow_all_targets, or pick a different service." - ) + unusable = _provider_service_unusable_reason(tool, match, service_name) + if unusable is not None: + return None, unusable return match, None @@ -3440,6 +3470,9 @@ def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | return catalog.model_ids, catalog.error_msg +_NO_ANTHROPIC_MODEL_IDS = "AI Gateway returned no Anthropic model ids" + + def list_anthropic_model_catalog( workspace: str, token: str, @@ -3484,7 +3517,7 @@ def list_anthropic_model_catalog( return AnthropicModelCatalog( model_ids=[], model_id_to_display_name={}, - error_msg="AI Gateway returned no Anthropic model ids", + error_msg=_NO_ANTHROPIC_MODEL_IDS, ) @@ -3532,6 +3565,27 @@ def fetch_ai_gateway_claude_models(workspace: str, token: str) -> dict[str, str] return models +# Drops a `..` or `.anthropic.` prefix, a region-less `anthropic.`, +# `databricks-` and a trailing `[...]` hint: +# system.ai.claude-sonnet-4-6 -> claude-sonnet-4-6 +# us.anthropic.claude-sonnet-4-6 -> claude-sonnet-4-6 +# anthropic.claude-opus-4-8 -> claude-opus-4-8 +# databricks-claude-sonnet-4-6 -> claude-sonnet-4-6 +# system.ai.claude-opus-4-7[1m] -> claude-opus-4-7 +# main.ucode.claude_haiku -> claude_haiku +# gpt-5.5 -> gpt-5.5 (one dot is a version, not a schema) +_BARE_MODEL_NAME_PATTERN = re.compile( + r"^(?:[^.\[]+\.[^.\[]+\.)?(?:anthropic\.)?(?:databricks-)?(?P[^\[]*)" +) + + +def bare_model_name(model: str) -> str: + """``model`` lowercased, without its catalog/schema or region prefix, ``anthropic.`` and + ``databricks-`` prefixes, or ``[...]`` hint.""" + match = _BARE_MODEL_NAME_PATTERN.match(model.lower()) + return match.group("name") if match else "" + + def model_version_sort_key(name: str) -> tuple: """Sort endpoint names so newer model versions come first. @@ -3619,7 +3673,7 @@ def discover_gemini_models(workspace: str, token: str) -> tuple[list[str], str | # Order newest model version first so `default_model()` (which picks the # first entry) launches e.g. gemini-3.5-flash rather than gemini-2.5-flash. return discover_endpoints_with_api_type( - workspace, token, "gemini/v1/generateContent", sort_key=model_version_sort_key + workspace, token, API_GEMINI_GENERATE_CONTENT, sort_key=model_version_sort_key ) @@ -3627,7 +3681,7 @@ def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str | # Order newest model version first (like `discover_gemini_models`), so the picker's top choice # and default is e.g. gpt-5-4 rather than the alphabetically-first gpt-5. return discover_endpoints_with_api_type( - workspace, token, "openai/v1/responses", sort_key=model_version_sort_key + workspace, token, API_OPENAI_RESPONSES, sort_key=model_version_sort_key ) @@ -3862,10 +3916,33 @@ def _fetch_codex_model_catalog( if not isinstance(payload, dict) or not isinstance(payload.get("models"), list): raise RuntimeError(f"{kind} {identifier} returned an invalid Codex model catalog.") if not payload["models"]: - raise RuntimeError(f"{kind} {identifier} returned no Codex models.") + raise CodexModelCatalogEmpty(f"{kind} {identifier} returned no Codex models.") return payload +def list_provider_service_models( + workspace: str, token: str, service_name: str, provider_type: str +) -> tuple[list[str], str | None]: + """The models the gateway lists for Model Provider Service ``service_name``, and why listing + failed. An empty listing is not a failure. + + Anthropic and Bedrock services list through the Anthropic models route, others through the + Codex model catalog.""" + if provider_type in ANTHROPIC_API_PROVIDER_TYPES: + listing = list_anthropic_model_catalog(workspace, token, provider=service_name) + reason = listing.error_msg + return listing.model_ids, None if reason == _NO_ANTHROPIC_MODEL_IDS else reason + try: + catalog = _fetch_codex_model_catalog( + workspace, token, source=CodexCatalogSource.PROVIDER, identifier=service_name + ) + except CodexModelCatalogEmpty: + return [], None + except RuntimeError as exc: + return [], str(exc) + return catalog_slugs(catalog), None + + def build_opencode_base_urls(workspace: str) -> dict[str, str]: return { "anthropic": build_tool_base_url("claude", workspace) + "/v1", diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index b2c89f5b9..6e8c43f3a 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -34,7 +34,7 @@ import threading import time import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import cast from urllib.parse import unquote @@ -270,6 +270,10 @@ 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 _request_body(self, body: bytes | None) -> bytes | None: + """The body to forward; ``body`` unchanged unless a subclass rewrites it.""" + return body + def _added_headers(self) -> dict[str, str]: return {} @@ -282,7 +286,7 @@ def _handle(self) -> None: diagnostic_id = uuid.uuid4().hex[:12] started = time.monotonic() length = int(self.headers.get("Content-Length", 0) or 0) - body = self.rfile.read(length) if length else None + body = self._request_body(self.rfile.read(length) if length else None) url = self.path.lstrip("/") token_header, extra_strip, route_label = self._forward_target(body) log_proxy_diagnostic( @@ -487,6 +491,24 @@ class _GatewayAuthProxyHandler(_ProxyHandler): added_headers: Mapping[str, str] allowed_path_prefixes: tuple[str, ...] sign_in_hint: str + dropped_chat_completions_fields: frozenset[str] + + def _request_body(self, body: bytes | None) -> bytes | None: + if not body or not self.dropped_chat_completions_fields: + return body + if not self.path.split("?", 1)[0].endswith("/chat/completions"): + return body + try: + payload = json.loads(body) + except ValueError: + return body + if not isinstance(payload, dict) or self.dropped_chat_completions_fields.isdisjoint( + payload + ): + return body + for field in self.dropped_chat_completions_fields: + payload.pop(field, None) + return json.dumps(payload).encode() def _has_session_secret(self) -> bool: expected = self.session_secret.encode() @@ -658,12 +680,15 @@ def start_gateway_auth_proxy( added_headers: Mapping[str, str], sign_in_hint: str, on_refresh_failure: Callable[[BaseException], None] = log_token_refresh_failure, + dropped_chat_completions_fields: Iterable[str] = (), ) -> 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]``. + ``added_headers``, after removing ``dropped_chat_completions_fields`` from Chat + Completions request bodies. Binds an OS-assigned port, read from + ``server.server_address[1]``. """ return _start_proxy( workspace, @@ -679,6 +704,7 @@ def start_gateway_auth_proxy( "added_headers": dict(added_headers), "allowed_path_prefixes": GATEWAY_MODEL_PATH_PREFIXES, "sign_in_hint": sign_in_hint, + "dropped_chat_completions_fields": frozenset(dropped_chat_completions_fields), }, on_refresh_failure=on_refresh_failure, ) diff --git a/tests/integration/README.md b/tests/integration/README.md index 9d9dde8e9..c581f6afd 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -270,6 +270,9 @@ and then a GPT task chosen with Copilot's `dbx-responses/` form, an unkno fails before Copilot starts, and the TUI switches from Claude to GPT with `/model` mid-session. Model evidence comes from Copilot's own `~/.copilot/session-state/*/events.jsonl` (`utils/agents/copilot.py`). +`test_ug_copilot_providers.py` runs Copilot through the MPS fixtures below: a file task on the +Anthropic service's model, on the OpenAI and Azure OpenAI services' `gpt-5-nano`, and on a named +model of the Bedrock allow-all service; launching on the relayed service fails before Copilot starts. ## Shared agent helpers @@ -496,6 +499,7 @@ test_ug_codex_headless.py # script prompts and model arguments test_ug_copilot_headless.py # Copilot via the session proxy: Claude + GPT tasks, no token on disk test_ug_copilot_lifecycle.py # Copilot refresh after revocation, sign-in expired, configure/revert, upgrade test_ug_copilot_models.py # Copilot model registry: Claude then GPT, unknown model, TUI /model switch +test_ug_copilot_providers.py # Copilot via Model Provider Services: Anthropic, OpenAI, Azure OpenAI, Bedrock allow-all; relayed refused 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 @@ -626,6 +630,7 @@ MPS CUJs select the existing services already used by e2e: requires `CLAUDE_CODE_OAUTH_TOKEN` (from `claude setup-token`) so the relayed launch runs headless instead of opening a browser login. - Codex: `main.ucode.ci_openai_mps`, using its allowed `gpt-5-nano` model. +- Codex and Copilot (Azure OpenAI): `main.ucode.ci_azure_openai_mps`, using its `gpt-5-nano` model. Use `--claude-provider` / `--claude-relayed-provider` / `--codex-provider` to reproduce another existing service. Use `--claude-provider-model` / @@ -697,7 +702,7 @@ cases instead complete real file tasks through `--workspace` and `--model-locati journeys check setup and launch with existing services, without inference. There are **87 live cases** (including 11 marked TUI journeys) and **7 installation -checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds ten live cases. One +checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds fifteen 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** (twelve of them also `live`) run on the managed workspace with a checked-in JSON CodingAgentConfig from `tests/fixtures/managed_config/` @@ -923,7 +928,7 @@ A second non-blocking job, **Managed OpenCode self-managed** (`managed_fixture a in `test_ug_agents_self_managed.py`), installs Claude, Codex and OpenCode and runs against the managed e2e workspace (`E2E_ADMIN_WORKSPACE`) with the injected `managed_workspace_default` config, which enables Claude/Codex but not OpenCode. It is likewise outside `cujs`. -A non-blocking **Copilot** job (`live and copilot`, nine cases) runs on Linux and Windows the +A non-blocking **Copilot** job (`live and copilot`, fourteen cases) runs on Linux and Windows the same way and is likewise outside `cujs` until it is stable. Each lane installs only the agent CLIs it needs (one, except the managed OpenCode lane's diff --git a/tests/integration/test_ug_copilot_providers.py b/tests/integration/test_ug_copilot_providers.py new file mode 100644 index 000000000..01067331d --- /dev/null +++ b/tests/integration/test_ug_copilot_providers.py @@ -0,0 +1,150 @@ +"""CUJs for running GitHub Copilot CLI through a Unity Catalog Model Provider Service via ug.""" + +import pytest +from utils.constants import COPILOT_HEADLESS_FLAGS +from utils.evidence import ( + FileTask, + assert_completed_task_model, + assert_no_session_registry_left, + copilot_selected_models, +) + +pytestmark = [pytest.mark.live, pytest.mark.copilot] + + +def _configure(session, workspace): + session.run( + "configure", + "--agents", + "copilot", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + + +def _run_task(session, provider, *model_args): + task = FileTask(session) + result = session.run( + "copilot", + "--provider", + provider, + *model_args, + "-p", + task.prompt, + *COPILOT_HEADLESS_FLAGS, + timeout=180, + ) + return task, result + + +def test_ug_copilot_runs_a_task_through_an_anthropic_provider( + live_session, workspace, claude_provider, claude_provider_model +): + """Scenario: configure Copilot, then run a headless task with `--provider` set to an API-key + Anthropic provider service and no `--model`. + + Expected: ug names the provider, Copilot starts on the service's Claude model through + `dbx-anthropic`, the completed answer contains the fixture value and was sent as that model, + and no session registry is left. + """ + session = live_session + _configure(session, workspace) + session.record("copilot-provider.json", {"provider": claude_provider}) + + task, result = _run_task(session, claude_provider) + + assert claude_provider in result.stderr + result.stdout + task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, claude_provider_model) + assert f"dbx-anthropic/{claude_provider_model}" in copilot_selected_models(session) + assert_no_session_registry_left(session) + + +def test_ug_copilot_runs_a_task_through_an_openai_provider( + live_session, workspace, codex_provider, codex_provider_model +): + """Scenario: configure Copilot, then run a headless task through an OpenAI provider service + on its allowed GPT model. + + Expected: the session runs on `dbx-responses/`, the completed answer contains the + fixture value and was sent as that model, and no session registry is left. + """ + session = live_session + _configure(session, workspace) + session.record("copilot-provider.json", {"provider": codex_provider}) + + task, result = _run_task(session, codex_provider, "--model", codex_provider_model) + + task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, codex_provider_model) + assert f"dbx-responses/{codex_provider_model}" in copilot_selected_models(session) + assert_no_session_registry_left(session) + + +def test_ug_copilot_runs_a_task_through_an_azure_openai_provider( + live_session, workspace, codex_azure_provider, codex_azure_provider_model +): + """Scenario: configure Copilot, then run a headless task through an Azure OpenAI provider + service on its GPT model. + + Expected: the session runs on `dbx-responses/` through the gateway's native OpenAI + route, the completed answer contains the fixture value and was sent as that model, and no + session registry is left. + """ + session = live_session + _configure(session, workspace) + session.record("copilot-provider.json", {"provider": codex_azure_provider}) + + task, result = _run_task(session, codex_azure_provider, "--model", codex_azure_provider_model) + + task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, codex_azure_provider_model) + assert f"dbx-responses/{codex_azure_provider_model}" in copilot_selected_models(session) + assert_no_session_registry_left(session) + + +def test_ug_copilot_allow_all_bedrock_provider_runs_the_named_model( + live_session, workspace, claude_bedrock_allow_all_provider, claude_bedrock_allow_all_model +): + """Scenario: run a headless task through a Bedrock provider service that allows every model + and declares none, naming its Bedrock Claude id with `--model`. + + Expected: ug accepts the provider, the completed answer contains the fixture value and was + sent as the Bedrock id, and no session registry is left. + """ + session = live_session + _configure(session, workspace) + session.record("copilot-provider.json", {"provider": claude_bedrock_allow_all_provider}) + + task, result = _run_task( + session, claude_bedrock_allow_all_provider, "--model", claude_bedrock_allow_all_model + ) + + task.assert_headless_answer("copilot", result) + assert_completed_task_model(session, "copilot", task.value, claude_bedrock_allow_all_model) + assert_no_session_registry_left(session) + + +def test_ug_copilot_refuses_a_subscription_relay_provider( + live_session, workspace, claude_relayed_provider +): + """Scenario: configure Copilot, then launch it with a provider service that relays a Claude + subscription. + + Expected: ug exits 1 before starting Copilot, and its error says only Claude Code can use + that service. No session registry is left. + """ + session = live_session + _configure(session, workspace) + + result = session.run( + "copilot", "--provider", claude_relayed_provider, "-p", "hi", ok=False, timeout=120 + ) + + assert result.returncode == 1, result.stdout + result.stderr + assert "only Claude Code can use" in result.stderr + assert not list(session.home.glob(".copilot/session-state/*")), "Copilot started" + assert_no_session_registry_left(session) diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 83d3eeb39..fe624e1d9 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -309,6 +309,12 @@ class TestCopilotCatalogModelId: ("system.ai.claude-haiku-4-5", "claude-haiku-4.5"), ("system.ai.claude-haiku-5-5", "claude-sonnet-5.5"), ("databricks-claude-sonnet-4-6", "claude-sonnet-4.6"), + ("claude-sonnet-4-5-20250929", "claude-sonnet-4.5"), + ("us.anthropic.claude-sonnet-4-6", "claude-sonnet-4.6"), + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "claude-sonnet-4.5"), + ("global.anthropic.claude-opus-4-8", "claude-opus-4.8"), + ("anthropic.claude-opus-4-8", "claude-opus-4.8"), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", "claude-sonnet-4.5"), ], ) def test_maps_claude_ids(self, model, expected): @@ -318,6 +324,220 @@ def test_other_models_have_no_catalog_id(self): assert copilot.copilot_catalog_model_id(GPT) is None +class TestClassifyProviderModel: + @pytest.mark.parametrize( + ("model", "provider_type", "wire"), + [ + ("claude-sonnet-5", "anthropic", "anthropic"), + ("us.anthropic.claude-sonnet-4-6", "amazon_bedrock", "anthropic"), + ("anthropic.claude-opus-4-8", "bedrock_mantle", "anthropic"), + ("gpt-5.5", "openai", "responses"), + ("gpt-5-nano", "openai", "responses"), + ("o3-mini", "openai", "responses"), + ("gpt-4o", "azure_openai", "completions"), + ("my-gpt4o-deployment", "azure_openai", "completions"), + ], + ) + def test_provider_type_and_id_decide_the_wire(self, model, provider_type, wire): + assert copilot.classify_model(model, provider_type=provider_type) == (wire, None) + + def test_declared_non_chat_apis_exclude(self): + wire, why = copilot.classify_model( + "embed-v-4-0", ["openai/v1/embeddings"], provider_type="openai" + ) + + assert wire is None + assert "isn't a chat model" in why + assert copilot.classify_model( + "gpt-5.2-chat", ["openai/v1/chat/completions"], provider_type="azure_openai" + ) == ("responses", None) + + @pytest.mark.parametrize( + ("model", "provider_type", "reason"), + [ + ("us.anthropic.claude-sonnet-4-6", "openai", None), + ("amazon.titan-text-express-v1", "amazon_bedrock", "Claude models"), + ], + ) + def test_unusable_models_say_why(self, model, provider_type, reason): + wire, why = copilot.classify_model(model, provider_type=provider_type) + if reason is None: + assert wire == "completions" + else: + assert wire is None + assert reason in why + + +def provider_models(models, provider_type, api_types=None): + api_types = api_types or {} + return copilot.build_catalog( + ((model, api_types.get(model, ())) for model in models), provider_type=provider_type + ) + + +class TestBuildProviderCatalog: + def test_keeps_usable_models_in_catalog_order_and_reports_the_rest(self): + models, skipped = provider_models( + ["gpt-4o", "text-embedding-3-small", "gpt-5.5", "gpt-5.5"], + "openai", + {"text-embedding-3-small": ["openai/v1/embeddings"]}, + ) + + assert [(m.id, m.wire) for m in models] == [ + ("gpt-5.5", "responses"), + ("gpt-4o", "completions"), + ] + assert skipped == [ + ("text-embedding-3-small", "it isn't a chat model (it serves openai/v1/embeddings)") + ] + + def test_orders_dotted_provider_versions_newest_first(self): + models, _ = provider_models( + ["gpt-5", "gpt-5.6", "gpt-5.1", "gpt-5.5", "o1", "o4-mini", "o3"], "openai" + ) + + assert [m.id for m in models] == [ + "gpt-5.6", + "gpt-5.5", + "gpt-5.1", + "gpt-5", + "o4-mini", + "o3", + "o1", + ] + + +def _stub_provider_service(monkeypatch, service, *, listed=(), listing_error=None): + monkeypatch.setattr(copilot, "launch_token", lambda state, workspace, **_: "dbx-token") + seen = {} + + def resolve(tool, name, workspace, token): + seen["resolve"] = (tool, name, workspace, token) + return (dict(service, name=name), None) if service else (None, f"'{name}' was not found.") + + def list_models(workspace, token, service_name, provider_type): + seen["listing"] = (service_name, provider_type) + return list(listed), listing_error + + monkeypatch.setattr(copilot, "resolve_provider_service", resolve) + monkeypatch.setattr(copilot, "list_provider_service_models", list_models) + return seen + + +MPS = "main.default.anthropic_mps" + + +class TestProviderCatalog: + def test_uses_the_declared_targets(self, monkeypatch): + seen = _stub_provider_service( + monkeypatch, + {"provider_type": "anthropic", "targets": ["claude-sonnet-5", "claude-opus-5"]}, + ) + + models = copilot.provider_service_catalog({"workspace": WS}, MPS, None) + + assert [m.id for m in models] == ["claude-opus-5", "claude-sonnet-5"] + assert seen["resolve"] == ("copilot", MPS, WS, "dbx-token") + assert "anthropic_listing" not in seen + + def test_allow_all_adds_the_gateway_listing_to_the_start_model(self, monkeypatch): + seen = _stub_provider_service( + monkeypatch, + {"provider_type": "amazon_bedrock", "targets": [], "allow_all_targets": True}, + listed=["anthropic.claude-opus-4-8"], + ) + + models = copilot.provider_service_catalog( + {"workspace": WS}, MPS, "us.anthropic.claude-sonnet-4-6" + ) + + assert [m.id for m in models] == [ + "anthropic.claude-opus-4-8", + "us.anthropic.claude-sonnet-4-6", + ] + assert seen["listing"] == (MPS, "amazon_bedrock") + + def test_allow_all_openai_takes_a_provider_ref_as_the_start_model(self, monkeypatch): + seen = _stub_provider_service( + monkeypatch, + {"provider_type": "openai", "targets": [], "allow_all_targets": True}, + listed=["gpt-5.5", "gpt-4o"], + ) + + models = copilot.provider_service_catalog({"workspace": WS}, MPS, "dbx-responses/gpt-5.5") + + assert [(m.id, m.wire) for m in models] == [ + ("gpt-5.5", "responses"), + ("gpt-4o", "completions"), + ] + assert seen["listing"] == (MPS, "openai") + + def test_allow_all_starts_on_the_requested_model_when_listing_fails(self, monkeypatch, capsys): + _stub_provider_service( + monkeypatch, + {"provider_type": "openai", "targets": [], "allow_all_targets": True}, + listing_error="HTTP 503 Service Unavailable", + ) + + models = copilot.provider_service_catalog({"workspace": WS}, MPS, "gpt-5.5") + + assert [m.id for m in models] == ["gpt-5.5"] + assert "HTTP 503" in capsys.readouterr().out + + def test_allow_all_without_a_start_model_registers_the_listing(self, monkeypatch): + _stub_provider_service( + monkeypatch, + {"provider_type": "anthropic", "targets": [], "allow_all_targets": True}, + listed=["claude-opus-5", "claude-sonnet-5"], + ) + + models = copilot.provider_service_catalog({"workspace": WS}, MPS, None) + + assert [m.id for m in models] == ["claude-opus-5", "claude-sonnet-5"] + assert copilot.launch_model({}, models, warn=False).id == "claude-sonnet-5" + + def test_allow_all_with_an_empty_listing_asks_for_a_model(self, monkeypatch): + _stub_provider_service( + monkeypatch, {"provider_type": "anthropic", "targets": [], "allow_all_targets": True} + ) + + with pytest.raises(RuntimeError, match="has no models ug copilot can use.*--model"): + copilot.provider_service_catalog({"workspace": WS}, MPS, None) + + def test_no_usable_model_lists_the_reasons(self, monkeypatch): + _stub_provider_service( + monkeypatch, + { + "provider_type": "openai", + "targets": ["text-embedding-3-large"], + "target_api_types": {"text-embedding-3-large": ("openai/v1/embeddings",)}, + }, + ) + + with pytest.raises(RuntimeError, match=r"text-embedding-3-large \(it isn't a chat model"): + copilot.provider_service_catalog({"workspace": WS}, MPS, None) + + def test_warns_about_skipped_models(self, monkeypatch, capsys): + _stub_provider_service( + monkeypatch, + { + "provider_type": "openai", + "targets": ["gpt-5.5", "text-embedding-3-large"], + "target_api_types": {"text-embedding-3-large": ("openai/v1/embeddings",)}, + }, + ) + + copilot.provider_service_catalog({"workspace": WS}, MPS, None) + + assert "text-embedding-3-large" in capsys.readouterr().out + + def test_invalid_service_raises_its_error(self, monkeypatch): + _stub_provider_service(monkeypatch, None) + + with pytest.raises(RuntimeError, match="was not found"): + copilot.provider_service_catalog({"workspace": WS}, MPS, None) + + class TestCopilotModelLimits: @pytest.mark.parametrize( ("model", "limits"), @@ -381,6 +601,57 @@ def test_models_carry_claude_settings_and_limits(self): {"id": GPT, "provider": "dbx-responses", "wireModel": GPT, "name": GPT}, ] + def test_provider_service_routes_every_provider_by_header(self): + claude, _ = provider_models(["claude-sonnet-4-6"], "anthropic") + gpt, _ = provider_models(["gpt-5.5", "gpt-4o"], "openai") + + config = copilot.render_providers_config( + claude + gpt, PROXY, "secret", provider_service=MPS + ) + + assert config["providers"] == [ + { + "name": "dbx-anthropic", + "type": "anthropic", + "baseUrl": f"{PROXY}/ai-gateway/anthropic", + "bearerToken": "secret", + "headers": {"Databricks-Model-Provider-Service": MPS}, + }, + { + "name": "dbx-responses", + "type": "openai", + "wireApi": "responses", + "baseUrl": f"{PROXY}/ai-gateway/openai/v1", + "bearerToken": "secret", + "headers": {"Databricks-Model-Provider-Service": MPS}, + }, + { + "name": "dbx-completions", + "type": "openai", + "wireApi": "completions", + "baseUrl": f"{PROXY}/ai-gateway/openai/v1", + "bearerToken": "secret", + "headers": {"Databricks-Model-Provider-Service": MPS}, + }, + ] + + def test_provider_service_claude_keeps_its_id_with_200k_limits(self): + claude, _ = provider_models(["claude-sonnet-4-6"], "anthropic") + + config = copilot.render_providers_config(claude, PROXY, "secret", provider_service=MPS) + + assert config["models"] == [ + { + "id": "claude-sonnet-4-6", + "provider": "dbx-anthropic", + "wireModel": "claude-sonnet-4-6", + "name": "claude-sonnet-4-6", + "modelId": "claude-sonnet-4.6", + "maxPromptTokens": 168_000, + "maxOutputTokens": 32_000, + } + ] + class TestBuildRuntimeEnv: def test_points_copilot_at_the_registry(self, tmp_path): @@ -896,6 +1167,39 @@ def popen(argv, *, env): with pytest.raises(httpx.ConnectError): httpx.post(f"{seen['base']}/v1/messages", content=b"{}", timeout=10) + def test_provider_service_launch_registers_its_models(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + seen = _stub_launch(monkeypatch) + _stub_provider_service( + monkeypatch, + {"provider_type": "anthropic", "targets": ["claude-opus-5", "claude-sonnet-5"]}, + ) + state = dict(STATE, copilot_static_models=[GPT]) + state[copilot.LAUNCH_PROVIDER_SERVICE_KEY] = MPS + + with pytest.raises(SystemExit): + copilot.launch(state, [], options=copilot.LaunchOptions()) + + registry = seen["registry"] + assert [m["id"] for m in registry["models"]] == ["claude-opus-5", "claude-sonnet-5"] + assert registry["providers"][0]["headers"] == {"Databricks-Model-Provider-Service": MPS} + assert seen["env"]["COPILOT_MODEL"] == "dbx-anthropic/claude-sonnet-5" + assert copilot.LAUNCH_PROVIDER_SERVICE_KEY not in state + + def test_provider_service_launch_rejects_a_model_it_does_not_serve(self, tmp_path, monkeypatch): + isolate_copilot_config_paths(monkeypatch, tmp_path) + seen = _stub_launch(monkeypatch) + _stub_provider_service( + monkeypatch, {"provider_type": "anthropic", "targets": ["claude-opus-5"]} + ) + state = dict(STATE) + state[copilot.LAUNCH_PROVIDER_SERVICE_KEY] = MPS + + with pytest.raises(RuntimeError, match="claude-opus-5"): + copilot.launch(state, ["--model", SONNET], options=copilot.LaunchOptions()) + + assert "argv" not in seen + def test_no_model_is_a_launch_error(self, tmp_path, monkeypatch): isolate_copilot_config_paths(monkeypatch, tmp_path) diff --git a/tests/test_cli.py b/tests/test_cli.py index 799275f88..a43af3589 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -24,6 +24,7 @@ import ucode.cli as cli_mod import ucode.databricks as db_mod import ucode.managed_files as managed_files +from ucode.agents import copilot as copilot_agent from ucode.cli import app from ucode.databricks import GatewayProbe from ucode.managed_config import ManagedConfigResult, normalize_managed_config @@ -2137,6 +2138,39 @@ def test_provider_sets_transient_codex_launch_marker(self): assert result.exit_code == 0, result.output assert mock_launch.call_args.args[1]["_codex_launch_provider"] == "main.default.openai" + def test_provider_sets_transient_copilot_launch_marker(self): + state = dict(MINIMAL_STATE, available_tools=["copilot"]) + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.resolve_provider_models", return_value=(None, None, False) + ) as mock_resolve, + patch("ucode.cli.configure_tool", return_value=state) as mock_configure, + patch("ucode.cli._fetch_managed_config", return_value=(None, False)), + patch("ucode.cli.set_current_workspace") as mock_set, + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke( + app, + [ + "copilot", + "--workspace", + "https://example.databricks.com/", + "--provider", + "main.default.openai", + ], + ) + + assert result.exit_code == 0, result.output + mock_set.assert_called_once_with("https://example.databricks.com") + assert mock_resolve.call_args.args[0] == "copilot" + assert mock_configure.call_args.args[2] is None + launch_state = mock_launch.call_args.args[1] + assert launch_state[copilot_agent.LAUNCH_PROVIDER_SERVICE_KEY] == "main.default.openai" + def test_model_location_sets_transient_codex_launch_marker(self): state = dict(MINIMAL_STATE) with ( diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 9910d1a1f..3ed2ddbe5 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -867,6 +867,7 @@ def test_strips_prefix_and_tags_provider_type(self, monkeypatch): "name": "main.schema1.anthropic-svc", "provider_type": "anthropic", "targets": [], + "target_api_types": {}, "allow_all_targets": False, "relayed": False, } @@ -896,6 +897,33 @@ def test_extracts_targets(self, monkeypatch): "global.anthropic.claude-opus-4-8", ] + def test_keeps_each_targets_declared_api_types(self, monkeypatch): + payload = { + "model_provider_services": [ + { + "name": "model-provider-services/main.schema1.foundry-svc", + "config": { + "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_MICROSOFT_FOUNDRY", + "targets": [ + {"model": "gpt-5.4-nano", "native_api_types": ["openai/v1/responses"]}, + {"model": "embed-v-4-0", "native_api_types": ["openai/v1/embeddings"]}, + {"model": "plain"}, + ], + }, + } + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + + services, _ = db_mod.list_model_provider_services(WS, "token") + + assert services[0]["target_api_types"] == { + "gpt-5.4-nano": ("openai/v1/responses",), + "embed-v-4-0": ("openai/v1/embeddings",), + } + def test_returns_reason_on_failure(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 500 Server Error") @@ -1239,6 +1267,43 @@ def test_wrong_type_rejected(self, monkeypatch): assert service is None assert "can't route to" in error + def test_copilot_accepts_openai(self, monkeypatch): + self._patch(monkeypatch) + service, error = db_mod.resolve_provider_service( + "copilot", "main.schema1.openai-svc", WS, "token" + ) + assert error is None + assert service["provider_type"] == "openai" + + def test_copilot_rejects_bedrock_without_claude(self, monkeypatch): + self._patch(monkeypatch) + service, error = db_mod.resolve_provider_service( + "copilot", "main.schema2.bedrock-titan-svc", WS, "token" + ) + assert service is None + assert "no Claude models" in error + + def test_copilot_rejects_a_subscription_relay(self, monkeypatch): + payload = { + "model_provider_services": [ + { + "name": "model-provider-services/main.schema1.relay-svc", + "config": { + "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_ANTHROPIC", + "anthropic": {"relayed": {}}, + }, + } + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + service, error = db_mod.resolve_provider_service( + "copilot", "main.schema1.relay-svc", WS, "token" + ) + assert service is None + assert "only Claude Code can use" in error + def test_bedrock_without_claude_rejected(self, monkeypatch): self._patch(monkeypatch) service, error = db_mod.resolve_provider_service( @@ -1361,6 +1426,60 @@ def test_gemini_enterprise_ok_for_gemini(self, monkeypatch): assert service["targets"] == ["gemini-3.5-flash"] +class TestListProviderServiceModels: + @pytest.mark.parametrize( + ("provider_type", "response", "expected"), + [ + ( + "amazon_bedrock", + ({"data": [{"id": "us.anthropic.claude-sonnet-4-6"}]}, None), + (["us.anthropic.claude-sonnet-4-6"], None), + ), + ("anthropic", ({"data": []}, None), ([], None)), + ("anthropic", (None, "HTTP 403 Forbidden"), ([], "HTTP 403 Forbidden")), + ("openai", ({"models": [{"slug": "gpt-5.5"}]}, None), (["gpt-5.5"], None)), + ("openai", ({"models": []}, None), ([], None)), + ], + ) + def test_reports_failures_but_not_empty_listings( + self, monkeypatch, provider_type, response, expected + ): + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token, **kw: response) + + assert ( + db_mod.list_provider_service_models(WS, "token", "main.a.b", provider_type) == expected + ) + + def test_reports_a_codex_listing_failure(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, **kw: (None, "HTTP 503 Service Unavailable"), + ) + + ids, reason = db_mod.list_provider_service_models(WS, "token", "main.a.b", "openai") + + assert ids == [] + assert "HTTP 503" in reason + + +@pytest.mark.parametrize( + ("model", "name"), + [ + ("system.ai.claude-sonnet-4-6", "claude-sonnet-4-6"), + ("System.AI.Claude-Opus-4-7[1m]", "claude-opus-4-7"), + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "claude-sonnet-4-5-20250929-v1:0"), + ("anthropic.claude-opus-4-8", "claude-opus-4-8"), + ("databricks-claude-sonnet-4-6", "claude-sonnet-4-6"), + ("main.ucode.claude_haiku", "claude_haiku"), + ("gpt-5.5", "gpt-5.5"), + ("main.default.gpt-5.5", "gpt-5.5"), + ], +) +def test_bare_model_name(model, name): + assert db_mod.bare_model_name(model) == name + + class TestServiceUsableForTool: def test_bedrock_allow_all_targets_usable_without_targets(self): service = {"provider_type": "amazon_bedrock", "targets": [], "allow_all_targets": True} @@ -1383,6 +1502,24 @@ def test_bedrock_claude_targets_required_only_for_claude(self, provider_type): } assert not db_mod.service_usable_for_tool("claude", service) assert db_mod.service_usable_for_tool("codex", service) + assert not db_mod.service_usable_for_tool("copilot", service) + + @pytest.mark.parametrize( + "provider_type", ["anthropic", "amazon_bedrock", "bedrock_mantle", "openai", "azure_openai"] + ) + def test_copilot_routes_anthropic_and_openai_services(self, provider_type): + service = {"provider_type": provider_type, "targets": [], "allow_all_targets": True} + assert db_mod.service_usable_for_tool("copilot", service) + + @pytest.mark.parametrize("provider_type", ["microsoft_foundry", "gemini_enterprise", "custom"]) + def test_copilot_skips_other_provider_types(self, provider_type): + service = {"provider_type": provider_type, "targets": [], "allow_all_targets": True} + assert not db_mod.service_usable_for_tool("copilot", service) + + def test_copilot_skips_a_subscription_relay(self): + service = {"provider_type": "anthropic", "targets": [], "relayed": True} + assert not db_mod.service_usable_for_tool("copilot", service) + assert db_mod.service_usable_for_tool("claude", service) class TestModelProviderFeatureUnavailable: diff --git a/tests/test_gateway_proxy_integration.py b/tests/test_gateway_proxy_integration.py index 3518f9865..fb91a3087 100644 --- a/tests/test_gateway_proxy_integration.py +++ b/tests/test_gateway_proxy_integration.py @@ -17,6 +17,7 @@ from __future__ import annotations import contextlib +import json import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -279,6 +280,7 @@ def _running_gateway_auth_proxy( *, added_headers=None, on_refresh_failure=None, + dropped_chat_completions_fields=(), ): server, cache, client = gateway_proxy.start_gateway_auth_proxy( gateway.base_url, @@ -287,6 +289,7 @@ def _running_gateway_auth_proxy( added_headers=added_headers or {}, sign_in_hint="Run `databricks auth login`.", on_refresh_failure=on_refresh_failure or (lambda exc: None), + dropped_chat_completions_fields=dropped_chat_completions_fields, ) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() @@ -329,6 +332,23 @@ def test_swaps_session_secret_for_fresh_token_and_adds_headers(self, make_gatewa assert req.header("X-Tag") == "t" assert req.body == b'{"model":"m"}' + def test_drops_configured_fields_from_chat_completions_only(self, make_gateway): + gw = make_gateway() + body = b'{"model":"gpt-4o-mini","reasoning_effort":"medium","messages":[]}' + with _running_gateway_auth_proxy( + gw, dropped_chat_completions_fields=("reasoning_effort",) + ) as proxy_url: + for path in ("openai/v1/chat/completions", "openai/v1/responses"): + httpx.post( + f"{proxy_url}/ai-gateway/{path}", + headers={"Authorization": f"Bearer {_SESSION_SECRET}"}, + content=body, + timeout=10, + ) + chat, responses = gw.requests + assert json.loads(chat.body) == {"model": "gpt-4o-mini", "messages": []} + assert responses.body == body + 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: diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index e4c9be432..c2a10c9a8 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -381,8 +381,13 @@ def test_codex_supports_openai_azure_foundry_and_bedrock(self): def test_claude_does_not_support_openai(self): assert not supports_provider_service("claude", "openai") + def test_copilot_supports_anthropic_openai_and_bedrock(self): + for provider_type in ("anthropic", "amazon_bedrock", "openai", "azure_openai"): + assert supports_provider_service("copilot", provider_type), provider_type + assert not supports_provider_service("copilot", "microsoft_foundry") + def test_other_agents_have_no_provider_support(self): - for tool in ("gemini", "opencode", "pi", "copilot"): + for tool in ("gemini", "opencode", "pi"): assert not supports_provider_service(tool, "anthropic"), tool