diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 614f6c7a4..ed979b368 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -213,7 +213,7 @@ def classify_model( 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 None, "only the Claude models of this service can be used" return wire, None if provider_type is not None: # Foundry serves its Claude targets on the Anthropic route but not on Chat Completions, @@ -245,16 +245,26 @@ def uses_responses_api(model: str) -> bool: return _wire_from_name(bare_model_name(model)) == WIRE_RESPONSES +def claude_model_name(model: str) -> str: + """A Claude model's own name, without its catalog/schema and release suffix: + ``us.anthropic.claude-sonnet-4-5-20250929-v1:0`` -> ``claude-sonnet-4-5``. + + Other models keep their bare name.""" + name = bare_model_name(model) + if not name.startswith("claude-"): + return name + return _CLAUDE_RELEASE_SUFFIX_PATTERN.sub("", name, count=1) + + 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 = bare_model_name(model) + tail = claude_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. @@ -279,6 +289,17 @@ def copilot_model_limits(model: str) -> tuple[int, int] | None: return _CLAUDE_1M_LIMITS if one_million else _CLAUDE_200K_LIMITS +def catalog_model_limits(model: CopilotModel, *, external: bool) -> tuple[int, int] | None: + """:func:`copilot_model_limits` for a Claude ``model`` of the catalog, or None for others. + + ``external`` (a Model Provider Service, or a location's model service) caps them at 200K, which + such a service may serve instead of the 1M window.""" + limits = copilot_model_limits(model.id) if model.wire == WIRE_ANTHROPIC else None + if limits is not None and external and sum(limits) > _CLAUDE_200K_WINDOW: + return _CLAUDE_200K_LIMITS + return limits + + @dataclass(frozen=True) class CopilotModel: """A Databricks model in Copilot's registry and the wire API it is served on.""" @@ -347,30 +368,36 @@ def _without_registry_prefix(model: str) -> str: def provider_service_catalog( - state: dict, service_name: str, requested: str | None + state: dict, service_name: str, requested: str | None, *, tool: str = "copilot" ) -> 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.""" + (the launch's ``--model``) or the managed default, so a model the listing omits can be named. + ``tool`` is ``vscode`` for Copilot Chat in VS Code, which has neither.""" workspace = state["workspace"] token = launch_token(state, workspace) - service, error = resolve_provider_service("copilot", service_name, workspace, token) + service, error = resolve_provider_service(tool, service_name, workspace, token) if error or service is None: raise RuntimeError(error or f"Model provider service '{service_name}' was not found.") + agent = "Copilot" if tool == "copilot" else "VS Code" 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") + managed_default = state.get("copilot_default_model") if tool == "copilot" else None 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 `." + f"Could not list the models of '{service_name}': {reason}." + + ( + " Name the one to start on with `ug copilot --model `." + if tool == "copilot" + else "" + ) ) if reason: print_warning(f"Could not list the models of {service_name}: {reason}.") @@ -386,12 +413,16 @@ def provider_service_catalog( 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 "") + f"Model provider service '{service_name}' has no models {agent} can use.{detail}" + + ( + " Name one it serves with `ug copilot --model `." + if allow_all and tool == "copilot" + else "" + ) ) if skipped: print_warning( - f"Skipping models of {service_name} that ug copilot can't use: " + f"Skipping models of {service_name} that {agent} can't use: " + "; ".join(f"{model} ({reason})" for model, reason in skipped) + "." ) @@ -649,14 +680,7 @@ def render_providers_config( 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 == WIRE_ANTHROPIC else None - if ( - limits is not None - and (provider_service or model_location) - and sum(limits) > _CLAUDE_200K_WINDOW - ): - # A Model Provider Service, or a location's model service, may not serve the 1M window. - limits = _CLAUDE_200K_LIMITS + limits = catalog_model_limits(model, external=bool(provider_service or model_location)) if limits is not None: entry["maxPromptTokens"], entry["maxOutputTokens"] = limits models.append(entry) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index bab6972e1..72c354e00 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -5,6 +5,9 @@ holds the workspace, its models and the auth command, and installs it with VS Code's `code` command. The extension reads that file from its own folder, so ug never edits the user's settings.json. `ug revert` uninstalls the extension. + +The models come from one source: the workspace's own (`system.ai`), a Model Provider Service +(`--provider`), or the model services of a Unity Catalog schema (`--model-location`). """ from __future__ import annotations @@ -21,9 +24,12 @@ from xml.sax.saxutils import escape from ucode.agents.copilot import ( - copilot_exclusion_reason, - copilot_model_limits, - uses_responses_api, + CopilotModel, + build_copilot_catalog, + catalog_model_limits, + claude_model_name, + location_catalog, + provider_service_catalog, ) from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command, launch_token from ucode.databricks import build_auth_shell_command, discover_model_services @@ -161,36 +167,57 @@ def install_extension(code: Path, config: dict[str, object]) -> str: return f"installed {EXTENSION_ID} {version}" -def _model_entry(model: str, name: str) -> dict[str, object]: +def _model_entry(model: CopilotModel, name: str, *, external: bool) -> dict[str, object]: entry: dict[str, object] = { - "id": model, + "id": model.id, "name": name, - "api": "responses" if uses_responses_api(model) else "chat_completions", + "family": claude_model_name(model.id), + # The wire names of Copilot's catalog; the extension dispatches on them (see APIS in + # vscode-extension/src/extension.ts). + "api": model.wire, } - # Claude's limits hold on Chat Completions too. Models without known limits get the - # extension's 128K-in, 32K-out default. - if limits := copilot_model_limits(model): + # Models without known limits get the extension's 128K-in, 32K-out default. + if limits := catalog_model_limits(model, external=external): entry["maxInputTokens"], entry["maxOutputTokens"] = limits return entry -def model_entries(workspace: str, token: str) -> tuple[list[dict[str, object]], str | None]: - """The workspace's Claude and GPT models, each with the API the extension calls it on.""" - # The models Copilot CLI offers: Gemini and open models, gpt-oss among them, fail on Copilot's - # requests or aren't served on every workspace. - claude, gpt, _gemini, _oss, reason = discover_model_services(workspace, token) - models = [ - model - for model in dict.fromkeys([*claude.values(), *gpt]) - if copilot_exclusion_reason(model) is None - ] - labels = {model: model_picker_label(model.rsplit(".", 1)[-1]) for model in models} +def _catalog(state: dict, provider: str | None, location: str | None) -> list[CopilotModel]: + """The models Copilot CLI would offer from the same source; raises when there are none.""" + if provider: + return provider_service_catalog(state, provider, None, tool="vscode") + if location: + return location_catalog(state, location) + workspace = state["workspace"] + # Gemini and open models, gpt-oss among them, fail on Copilot's requests or aren't served on + # every workspace. + claude, gpt, _gemini, _oss, reason = discover_model_services( + workspace, launch_token(state, workspace) + ) + catalog = build_copilot_catalog([*claude.values(), *gpt]) + if not catalog: + raise RuntimeError( + f"No Claude or GPT models for VS Code ({reason or 'the workspace has none'})." + ) + return catalog + + +def model_entries( + state: dict, *, provider: str | None = None, location: str | None = None +) -> list[dict[str, object]]: + """The models of the chosen source, each with the API the extension calls it on.""" + catalog = _catalog(state, provider, location) + if provider: + # A service's targets are the provider's own ids, which is what its users know it by, and + # several of them differ only by a region the label would drop. + return [_model_entry(model, model.id, external=True) for model in catalog] + labels = {model: model_picker_label(claude_model_name(model.id)) for model in catalog} counts = Counter(labels.values()) # As in Codex's picker: two models with the same label show their full ids instead. return [ - _model_entry(model, model if counts[label] > 1 else label) + _model_entry(model, model.id if counts[label] > 1 else label, external=bool(location)) for model, label in labels.items() - ], reason + ] def auth_command(state: dict) -> str: @@ -218,9 +245,10 @@ def next_steps() -> list[str]: ] -def configure(state: dict) -> bool: - """Install the extension with this workspace's setup. Returns False, after a warning, if any - step fails.""" +def configure(state: dict, *, provider: str | None = None, location: str | None = None) -> bool: + """Install the extension with this workspace's setup and the models of ``provider`` (a Model + Provider Service), ``location`` (a ``.``) or else ``system.ai``. Returns False, + after a warning, if any step fails.""" code = find_code() if code is None: print_warning( @@ -228,19 +256,24 @@ def configure(state: dict) -> bool: "\"Shell Command: Install 'code' command in PATH\"), then rerun `ug configure --agents vscode`." ) return False - workspace = state["workspace"] - models, reason = model_entries(workspace, launch_token(state, workspace)) - if not models: + try: + models = model_entries(state, provider=provider, location=location) + except RuntimeError as exc: # Keep what an earlier configure installed: a failed listing shouldn't wipe a working setup. - print_warning( - f"No Claude or GPT models for VS Code ({reason or 'the workspace has none'}). " - "VS Code was left unchanged." - ) + print_warning(f"{exc} VS Code was left unchanged.") return False - config = {"workspace": workspace, "authCommand": auth_command(state), "models": models} + config: dict[str, object] = { + "workspace": state["workspace"], + "authCommand": auth_command(state), + "models": models, + } + if provider: + config["provider"] = provider + source = f" from {provider or location}" if provider or location else "" try: print_kv( - "VS Code extension", f"{install_extension(code, config)} with {len(models)} models" + "VS Code extension", + f"{install_extension(code, config)} with {len(models)} models{source}", ) except RuntimeError as exc: print_warning(f"{exc} Rerun `ug configure --agents vscode` to try again.") diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 96e84088a..dc9479f3a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -4561,6 +4561,22 @@ def configure( "workspace has an admin-managed config.", ), ] = None, + provider: Annotated[ + str | None, + typer.Option( + "--provider", + help="With --agents vscode: give VS Code the models of this Model Provider Service " + "(`..`) instead of the Databricks-hosted ones.", + ), + ] = None, + model_location: Annotated[ + str | None, + typer.Option( + "--model-location", + help="With --agents vscode: give VS Code the model services in this Unity Catalog " + "schema (`.`) instead of the Databricks-hosted ones.", + ), + ] = None, mcp: Annotated[ str | None, typer.Option( @@ -4627,6 +4643,21 @@ def configure( install_databricks_cli() if agent is not None and agents is not None: raise RuntimeError("Use either --agent or --agents, not both.") + if provider is not None or model_location is not None: + requested_agents = [a.strip().lower() for a in (agents or "").split(",") if a.strip()] + # Both flags choose VS Code's models alone, so another model agent in the same run + # would silently keep the workspace's own. Cursor has no models to choose. + if "vscode" not in requested_agents or [ + a for a in requested_agents if a not in ("vscode", "cursor") + ]: + raise RuntimeError( + "--provider and --model-location work only with `--agents vscode`. " + "Configure the other agents in a separate run." + ) + if provider is not None and model_location is not None: + raise RuntimeError("--provider and --model-location cannot be used together.") + if model_location is not None and not is_valid_catalog_schema(model_location): + raise RuntimeError("--model-location must be `.`.") # --workspaces / --profiles are deprecated aliases of the singular flags. if workspace is not None and workspaces is not None: raise RuntimeError("Use either --workspace or --workspaces, not both.") @@ -4722,7 +4753,7 @@ def configure( state = load_state() managed, _ = _read_policy_for_other_agent(state) refuse_user_scope_over_managed_config(state.get("workspace"), managed) - vscode.configure(state) + vscode.configure(state, provider=provider, location=model_location) elif mcp is not None: # MCP-only: `--mcp` without --agent(s) (e.g. Cursor, which isn't a # model agent, or adding MCP servers to an already-configured setup). diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a0accb9be..0fc2b9363 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2803,9 +2803,11 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: "microsoft_foundry", ), } +# Copilot Chat in VS Code can use the same provider types as Copilot CLI. +_TOOL_PROVIDER_TYPES["vscode"] = _TOOL_PROVIDER_TYPES["copilot"] # Agents that can only use a Bedrock service's Claude targets. -_CLAUDE_ONLY_BEDROCK_TOOLS = ("claude", "copilot") +_CLAUDE_ONLY_BEDROCK_TOOLS = ("claude", "copilot", "vscode") # Provider types that expose Bedrock-style model ids (e.g. # `us.anthropic.claude-sonnet-4-6`, or region-less `anthropic.claude-*` on @@ -3605,8 +3607,9 @@ def fetch_ai_gateway_claude_models(workspace: str, token: str) -> dict[str, str] # 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) +# A region prefix is dotted on Bedrock (`us.anthropic.`) and slashed on Azure (`eu/gpt-5.6-luna`). _BARE_MODEL_NAME_PATTERN = re.compile( - r"^(?:[^.\[]+\.[^.\[]+\.)?(?:anthropic\.)?(?:databricks-)?(?P[^\[]*)" + r"^(?:[^.\[]+\.[^.\[]+\.)?(?:[^/.\[]+/)?(?:anthropic\.)?(?:databricks-)?(?P[^\[]*)" ) diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index c473cfd19..cf995e308 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,4 @@ -"use strict";var K=Object.create;var f=Object.defineProperty;var Y=Object.getOwnPropertyDescriptor;var z=Object.getOwnPropertyNames;var X=Object.getPrototypeOf,Q=Object.prototype.hasOwnProperty;var V=(e,t)=>{for(var n in t)f(e,n,{get:t[n],enumerable:!0})},k=(e,t,n,s)=>{if(t&&typeof t=="object"||typeof t=="function")for(let o of z(t))!Q.call(e,o)&&o!==n&&f(e,o,{get:()=>t[o],enumerable:!(s=Y(t,o))||s.enumerable});return e};var T=(e,t,n)=>(n=e!=null?K(X(e)):{},k(t||!e||!e.__esModule?f(n,"default",{value:e,enumerable:!0}):n,e)),Z=e=>k(f({},"__esModule",{value:!0}),e);var ge={};V(ge,{CONFIG_FILE:()=>U,DatabricksModels:()=>v,activate:()=>ue,readConfig:()=>B,userAgent:()=>H});module.exports=Z(ge);var G=require("node:fs"),F=require("node:path"),g=T(require("vscode"));var b="0.1.0";var _=require("node:child_process"),te=5*6e4,ne=60*6e4,oe=e=>new Promise((t,n)=>{(0,_.exec)(e,{timeout:6e4,windowsHide:!0},(s,o,a)=>{s?n(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(a).trim().slice(0,500)||s.message}`)):t(String(o).trim())})}),se=e=>{try{let{exp:t}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof t=="number")return t*1e3}catch{}return Date.now()+ne},h=class{constructor(t,n=()=>{}){this.command=t;this.onFetch=n}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-te>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(t){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==t?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(t){return this.pending??=oe(t).then(n=>{if(!n)throw new Error(`Couldn't get a Databricks token: \`${t}\` printed nothing.`);return this.cached={value:n,expiresAt:se(n)},this.onFetch(this.cached.expiresAt),n}).finally(()=>{this.pending=void 0}),this.pending}};var l=T(require("vscode")),S=3,re=1e3,y=e=>e instanceof l.LanguageModelTextPart?e.value:e instanceof l.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof l.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",R=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,L=e=>e.inputSchema??{type:"object",properties:{}},P=(e,t,n)=>{let s;try{s=n.trim()?JSON.parse(n):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${t} with arguments that aren't a JSON object: ${n}`);return new l.LanguageModelToolCallPart(e,t,s)},E=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),O=()=>new Error("The response ended before the model finished. Try again."),A=(e,t)=>e.tools.length?{tools:e.tools.map(t),...e.toolRequired?{tool_choice:"required"}:{}}:{},j=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let t=0,n=0;for(let s of e.content){let o=s instanceof l.LanguageModelToolResultPart?s.content:[s];for(let a of o)t+=a instanceof l.LanguageModelToolCallPart?a.name.length+JSON.stringify(a.input).length:y(a).length,n+=R(a)===void 0?0:1}return Math.ceil(t/4)+n*re},x=e=>e.map(R).filter(t=>t!==void 0),I=e=>({results:e.filter(t=>t instanceof l.LanguageModelToolResultPart),calls:e.filter(t=>t instanceof l.LanguageModelToolCallPart),text:e.map(y).join(""),images:x(e)});function N(e){let t=e.messages.flatMap(n=>{let{results:s,calls:o,text:a,images:c}=I(n.content);if(Number(n.role)===S)return[{role:"system",content:a}];if(n.role===l.LanguageModelChatMessageRole.Assistant){let r=o.map(u=>({id:u.callId,type:"function",function:{name:u.name,arguments:JSON.stringify(u.input)}}));return a||r.length?[{role:"assistant",content:a||null,...r.length?{tool_calls:r}:{}}]:[]}let i=[...a?[{type:"text",text:a}]:[],...[...s.flatMap(r=>x(r.content)),...c].map(r=>({type:"image_url",image_url:{url:r}}))];return[...s.map(r=>({role:"tool",tool_call_id:r.callId,content:r.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:t,stream:!0,...A(e,n=>({type:"function",function:{name:n.name,description:n.description,parameters:L(n)}}))}}var ae=e=>typeof e=="string"?e:(e??[]).filter(t=>t.type==="text").map(t=>t.text??"").join(""),ie=new Set(["stop","tool_calls","function_call"]);async function*D(e){let t=[],n;for await(let s of e){let{error:o}=s;if(o)throw new Error(typeof o=="string"?o:o.message??JSON.stringify(o));for(let a of s.choices??[]){let c=ae(a.delta?.content);c&&(yield new l.LanguageModelTextPart(c));for(let i of a.delta?.tool_calls??[]){let r=t[i.index??0]??={id:"",name:"",args:""};r.id=i.id||r.id,r.name=i.function?.name||r.name,r.args+=i.function?.arguments??""}if(a.finish_reason&&(n=a.finish_reason,!ie.has(n)))throw E(n)}}if(!n)throw O();for(let s of t.filter(Boolean))yield P(s.id,s.name,s.args)}var le=e=>{let t=e.map(y).join(""),n=x(e);return n.length?[...t?[{type:"input_text",text:t}]:[],...n.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:t};function $(e){let t=e.messages.flatMap(n=>{let{results:s,calls:o,text:a,images:c}=I(n.content);if(Number(n.role)===S)return[{role:"system",content:a}];if(n.role===l.LanguageModelChatMessageRole.Assistant)return[...a?[{role:"assistant",content:a}]:[],...o.map(r=>({type:"function_call",call_id:r.callId,name:r.name,arguments:JSON.stringify(r.input)}))];let i=[...a?[{type:"input_text",text:a}]:[],...c.map(r=>({type:"input_image",image_url:r,detail:"auto"}))];return[...s.map(r=>({type:"function_call_output",call_id:r.callId,output:le(r.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:t,stream:!0,store:!1,...A(e,n=>({type:"function",name:n.name,description:n.description,parameters:L(n),strict:!1}))}}async function*q(e){let t=[],n=!1;for await(let s of e){let o=s;if(o.type==="response.output_text.delta"&&o.delta)yield new l.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")t.push(P(o.item.call_id??"",o.item.name??"",o.item.arguments??""));else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw E(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw O();yield*t}async function*C(e,t,n,s,o,a){let c=m=>fetch(`${e}/ai-gateway/mlflow/v1/${t}`,{method:"POST",headers:{...o,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(n),signal:a}),i=await s.get(),r=await c(i);if(r.status===401&&(await r.body?.cancel(),i=await s.renew(i),r=await c(i)),r.status===401)throw await r.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry.");if(!r.ok||!r.body)throw new Error(`AI Gateway returned ${r.status}: ${(await r.text()).slice(0,300)}`);let u=new TextDecoder,p="";for await(let m of r.body){p+=u.decode(m,{stream:!0});let M=p.split(` -`);p=M.pop()??"";for(let W of M){let w=J(W);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let d=J(p+u.decode());d&&d!=="[DONE]"&&(yield JSON.parse(d))}var J=e=>e.startsWith("data:")?e.slice(5).trim():"";var ce=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(t=>{let n=e[t];return typeof n=="number"&&Number.isInteger(n)&&n>0?[[t,n]]:[]})),U="ug-config.json",B=e=>{let t={};try{let o=JSON.parse(e??"{}");o&&typeof o=="object"&&!Array.isArray(o)&&(t=o)}catch{}let n=o=>t[o],s=n("models");return{workspace:String(n("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(o=>typeof o?.id=="string"&&o.id!==""&&(o.api==="chat_completions"||o.api==="responses")).map(o=>({id:o.id,...typeof o.name=="string"&&o.name.trim()?{name:o.name.trim()}:{},api:o.api,...ce(o)})),authCommand:String(n("authCommand")??"")}},H=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(t,n,s=()=>({})){this.settings=t;this.log=n;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:t,models:n,authCommand:s}=this.settings();return!t||!s||n.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):n.map(({id:o,name:a,api:c,maxInputTokens:i,maxOutputTokens:r})=>({id:o,api:c,name:a??o,family:o.split(".").pop()??o,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:r??32e3,capabilities:{toolCalling:!0,imageInput:c==="responses"||o.includes("claude")}}))}async provideLanguageModelChatResponse(t,n,s,o,a){let{workspace:c,authCommand:i}=this.settings();this.tokens??=new h(i,d=>this.log(`Fetched a Databricks token (expires ${new Date(d).toISOString()})`));let r={model:t.id,messages:n,tools:s.tools??[],toolRequired:s.toolMode===g.LanguageModelChatToolMode.Required},u=new AbortController,p=a.onCancellationRequested(()=>u.abort());try{let d=t.api==="responses"?q(C(c,"responses",$(r),this.tokens,this.headers(),u.signal)):D(C(c,"chat/completions",N(r),this.tokens,this.headers(),u.signal));for await(let m of d)o.report(m)}catch(d){if(!a.isCancellationRequested)throw this.log(`${t.id} failed: ${d instanceof Error?d.message:String(d)}`),d}finally{p.dispose()}}async provideTokenCount(t,n){return j(n)}},de=e=>{try{return(0,G.readFileSync)(e,"utf8")}catch{return}};function ue(e){let t=g.window.createOutputChannel("Unity Gateway",{log:!0}),n=B(de((0,F.join)(e.extensionPath,U))),s=()=>({"User-Agent":H(g.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),o=new v(()=>n,a=>t.info(a),s);e.subscriptions.push(t,g.lm.registerLanguageModelChatProvider("databricks-unity-gateway",o))}0&&(module.exports={CONFIG_FILE,DatabricksModels,activate,readConfig,userAgent}); +"use strict";var se=Object.create;var b=Object.defineProperty;var ae=Object.getOwnPropertyDescriptor;var re=Object.getOwnPropertyNames;var ie=Object.getPrototypeOf,ce=Object.prototype.hasOwnProperty;var le=(e,n)=>{for(var o in n)b(e,o,{get:n[o],enumerable:!0})},j=(e,n,o,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let t of re(n))!ce.call(e,t)&&t!==o&&b(e,t,{get:()=>n[t],enumerable:!(s=ae(n,t))||s.enumerable});return e};var L=(e,n,o)=>(o=e!=null?se(ie(e)):{},j(n||!e||!e.__esModule?b(o,"default",{value:e,enumerable:!0}):o,e)),de=e=>j(b({},"__esModule",{value:!0}),e);var Te={};le(Te,{CONFIG_FILE:()=>Q,DatabricksModels:()=>k,PROVIDER_HEADER:()=>X,activate:()=>Ce,readConfig:()=>Z,userAgent:()=>ee});module.exports=de(Te);var z=require("node:fs"),V=require("node:path"),h=L(require("vscode"));var M="0.1.0";var I=require("node:child_process"),pe=5*6e4,ue=60*6e4,me=e=>new Promise((n,o)=>{(0,I.exec)(e,{timeout:6e4,windowsHide:!0},(s,t,r)=>{s?o(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):n(String(t).trim())})}),fe=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+ue},x=class{constructor(n,o=()=>{}){this.command=n;this.onFetch=o}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-pe>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=me(n).then(o=>{if(!o)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:o,expiresAt:fe(o)},this.onFetch(this.cached.expiresAt),o}).finally(()=>{this.pending=void 0}),this.pending}};var l=L(require("vscode")),C=3,he=1e3,w=e=>e instanceof l.LanguageModelTextPart?e.value:e instanceof l.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof l.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",D=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,$=e=>{let n=D(e);return n&&`data:${n.mimeType};base64,${n.data}`},T=e=>e.inputSchema??{type:"object",properties:{}},R=(e,n,o)=>{let s;try{s=o.trim()?JSON.parse(o):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${o}`);return new l.LanguageModelToolCallPart(e,n,s)},S=e=>e==="max_output_tokens"||e==="length"||e==="max_tokens"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"||e==="refusal"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),P=()=>new Error("The response ended before the model finished. Try again."),q=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},G=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,o=0;for(let s of e.content){let t=s instanceof l.LanguageModelToolResultPart?s.content:[s];for(let r of t)n+=r instanceof l.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:w(r).length,o+=$(r)===void 0?0:1}return Math.ceil(n/4)+o*he},A=e=>e.map($).filter(n=>n!==void 0),E=e=>({results:e.filter(n=>n instanceof l.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof l.LanguageModelToolCallPart),text:e.map(w).join(""),images:A(e)});function J(e){let n=e.messages.flatMap(o=>{let{results:s,calls:t,text:r,images:d}=E(o.content);if(Number(o.role)===C)return[{role:"system",content:r}];if(o.role===l.LanguageModelChatMessageRole.Assistant){let a=t.map(i=>({id:i.callId,type:"function",function:{name:i.name,arguments:JSON.stringify(i.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let c=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>A(a.content)),...d].map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(w).join("")})),...c.length?[{role:"user",content:c}]:[]]});return{model:e.model,messages:n,stream:!0,...q(e,o=>({type:"function",function:{name:o.name,description:o.description,parameters:T(o)}}))}}var ye=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),ve=new Set(["stop","tool_calls","function_call"]);async function*F(e){let n=[],o;for await(let s of e){let{error:t}=s;if(t)throw new Error(typeof t=="string"?t:t.message??JSON.stringify(t));for(let r of s.choices??[]){let d=ye(r.delta?.content);d&&(yield new l.LanguageModelTextPart(d));for(let c of r.delta?.tool_calls??[]){let a=n[c.index??0]??={id:"",name:"",args:""};a.id=c.id||a.id,a.name=c.function?.name||a.name,a.args+=c.function?.arguments??""}if(r.finish_reason&&(o=r.finish_reason,!ve.has(o)))throw S(o)}}if(!o)throw P();for(let s of n.filter(Boolean))yield R(s.id,s.name,s.args)}var we=e=>{let n=e.map(w).join(""),o=A(e);return o.length?[...n?[{type:"input_text",text:n}]:[],...o.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function H(e){let n=[],o=e.messages.flatMap(s=>{let{results:t,calls:r,text:d,images:c}=E(s.content);if(Number(s.role)===C)return n.push(d),[];if(s.role===l.LanguageModelChatMessageRole.Assistant)return[...d?[{role:"assistant",content:d}]:[],...r.map(i=>({type:"function_call",call_id:i.callId,name:i.name,arguments:JSON.stringify(i.input)}))];let a=[...d?[{type:"input_text",text:d}]:[],...c.map(i=>({type:"input_image",image_url:i,detail:"auto"}))];return[...t.map(i=>({type:"function_call_output",call_id:i.callId,output:we(i.content)})),...a.length?[{role:"user",content:a}]:[]]});return{model:e.model,...n.length?{instructions:n.join(` + +`)}:{},input:o,stream:!0,store:!1,...q(e,s=>({type:"function",name:s.name,description:s.description,parameters:T(s),strict:!1}))}}async function*U(e){let n=[],o=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new l.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")n.push(R(t.item.call_id??"",t.item.name??"",t.item.arguments??""));else if(t.type==="response.completed")o=!0;else{if(t.type==="response.incomplete")throw S(t.response?.incomplete_details?.reason);if(t.type==="response.failed"||t.type==="error")throw new Error(t.response?.error?.message??t.message??"The model response failed.")}}if(!o)throw P();yield*n}var _=e=>e?[{type:"text",text:e}]:[],O=e=>e.flatMap(n=>{let o=D(n);return o?[{type:"image",source:{type:"base64",media_type:o.mimeType,data:o.data}}]:[]}),be={cache_control:{type:"ephemeral"}},N=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...be})};function B(e,n){let o=[],s=[];for(let t of e.messages){let{results:r,calls:d,text:c}=E(t.content);if(Number(t.role)===C){o.push(..._(c));continue}let a=t.role===l.LanguageModelChatMessageRole.Assistant?"assistant":"user",i=a==="assistant"?[..._(c),...d.map(g=>({type:"tool_use",id:g.callId,name:g.name,input:g.input}))]:[...r.map(g=>{let p=[..._(g.content.map(w).join("")),...O(g.content)];return{type:"tool_result",tool_use_id:g.callId,...p.length?{content:p}:{}}}),..._(c),...O(t.content)],u=s[s.length-1];if(i.length)u?.role===a?u.content.push(...i):s.push({role:a,content:i});else continue}return N(o),N(s[s.length-1]?.content??[]),{model:e.model,max_tokens:n,...o.length?{system:o}:{},messages:s,stream:!0,...e.tools.length?{tools:e.tools.map(t=>({name:t.name,description:t.description,input_schema:T(t)})),...e.toolRequired?{tool_choice:{type:"any"}}:{}}:{}}}var xe=new Set(["end_turn","tool_use","stop_sequence"]);async function*W(e){let n=new Map,o=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")n.set(t.index??0,{id:t.content_block.id??"",name:t.content_block.name??"",args:""});else if(t.type==="content_block_delta"&&t.delta?.type==="text_delta"&&t.delta.text)yield new l.LanguageModelTextPart(t.delta.text);else if(t.type==="content_block_delta"&&t.delta?.type==="input_json_delta"){let r=n.get(t.index??0);r&&(r.args+=t.delta.partial_json??"")}else{if(t.type==="message_delta"&&t.delta?.stop_reason&&!xe.has(t.delta.stop_reason))throw S(t.delta.stop_reason);if(t.type==="message_stop")o=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!o)throw P();for(let s of n.values())yield R(s.id,s.name,s.args)}async function*Y(e,n,o,s,t,r){let d=(p,f)=>fetch(`${e}${n}`,{method:"POST",headers:{...t,Authorization:`Bearer ${p}`,"Content-Type":"application/json"},body:JSON.stringify(f),signal:r}),c=await s.get(),a=await d(c,o);if(a.status===401&&(await a.body?.cancel(),c=await s.renew(c),a=await d(c,o)),a.status===401)throw await a.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry.");if(!a.ok||!a.body){let p=(await a.text()).slice(0,300),f=Number(/max_tokens: \d+ > (\d+)/.exec(p)?.[1]);if(!f)throw new Error(`AI Gateway returned ${a.status}: ${p}`);if(a=await d(c,{...o,max_tokens:f}),!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`)}let i=new TextDecoder,u="";for await(let p of a.body){u+=i.decode(p,{stream:!0});let f=u.split(` +`);u=f.pop()??"";for(let v of f){let y=K(v);if(y==="[DONE]")return;y&&(yield JSON.parse(y))}}let g=K(u+i.decode());g&&g!=="[DONE]"&&(yield JSON.parse(g))}var K=e=>e.startsWith("data:")?e.slice(5).trim():"";var _e=["anthropic","completions","responses"],X="Databricks-Model-Provider-Service",ke=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let o=e[n];return typeof o=="number"&&Number.isInteger(o)&&o>0?[[n,o]]:[]})),Q="ug-config.json",Z=e=>{let n={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(n=t)}catch{}let o=t=>n[t],s=o("models");return{workspace:String(o("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&_e.includes(t.api)).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},...typeof t.family=="string"&&t.family?{family:t.family}:{},api:t.api,...ke(t)})),authCommand:String(o("authCommand")??""),provider:typeof o("provider")=="string"?String(o("provider")):""}},ee=e=>`ucode/${M} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,k=class{constructor(n,o,s=()=>({})){this.settings=n;this.log=o;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:n,models:o,authCommand:s}=this.settings();return!n||!s||o.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):o.map(({id:t,name:r,family:d,api:c,maxInputTokens:a,maxOutputTokens:i})=>({id:t,api:c,name:r??t,family:d??t.split(".").pop()??t,version:M,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:i??32e3,capabilities:{toolCalling:!0,imageInput:c!=="completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(n,o,s,t,r){let{workspace:d,authCommand:c,provider:a}=this.settings();this.tokens??=new x(c,m=>this.log(`Fetched a Databricks token (expires ${new Date(m).toISOString()})`));let i={model:n.id,messages:o,tools:s.tools??[],toolRequired:s.toolMode===h.LanguageModelChatToolMode.Required},u=new AbortController,g=r.onCancellationRequested(()=>u.abort()),p={...this.headers(),...a?{[X]:a}:{}},f=this.tokens,v=(m,ne,oe={})=>Y(d,m,ne,f,{...p,...oe},u.signal),y=a?"/ai-gateway/openai/v1":"/ai-gateway/mlflow/v1",te={anthropic:()=>W(v("/ai-gateway/anthropic/v1/messages",B(i,n.maxOutputTokens),{"anthropic-version":"2023-06-01"})),responses:()=>U(v(`${y}/responses`,H(i))),completions:()=>F(v(`${y}/chat/completions`,J(i)))};try{for await(let m of te[n.api]())t.report(m)}catch(m){if(!r.isCancellationRequested)throw this.log(`${n.id} failed: ${m instanceof Error?m.message:String(m)}`),m}finally{g.dispose()}}async provideTokenCount(n,o){return G(o)}},Me=e=>{try{return(0,z.readFileSync)(e,"utf8")}catch{return}};function Ce(e){let n=h.window.createOutputChannel("Unity Gateway",{log:!0}),o=Z(Me((0,V.join)(e.extensionPath,Q))),s=()=>({"User-Agent":ee(h.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>o,r=>n.info(r),s);e.subscriptions.push(n,h.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,activate,readConfig,userAgent}); diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index b1ac9a0f0..2bebd3a3e 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -518,7 +518,7 @@ def test_allow_all_with_an_empty_listing_asks_for_a_model(self, monkeypatch): monkeypatch, {"provider_type": "anthropic", "targets": [], "allow_all_targets": True} ) - with pytest.raises(RuntimeError, match="has no models ug copilot can use.*--model"): + with pytest.raises(RuntimeError, match="has no models Copilot can use.*--model"): copilot.provider_service_catalog({"workspace": WS}, MPS, None) def test_no_usable_model_lists_the_reasons(self, monkeypatch): diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 48c55402a..27deea8cd 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -11,8 +11,10 @@ from typer.testing import CliRunner from ucode import config_io -from ucode.agents import vscode +from ucode import databricks as databricks_module +from ucode.agents import copilot, vscode from ucode.cli import app +from ucode.databricks import AnthropicModelCatalog from ucode.managed_files import OS # The conftest fixture stubs these out for every test; the ones here exercise the real lookups. @@ -65,6 +67,7 @@ def code(monkeypatch): monkeypatch.setattr(vscode, "_find_command", lambda command: Path(f"/fake/bin/{command}")) monkeypatch.setattr(vscode, "_run_code", fake) monkeypatch.setattr(vscode, "launch_token", lambda state, workspace: "token") + monkeypatch.setattr(copilot, "launch_token", lambda state, workspace: "token") monkeypatch.setattr( vscode, "discover_model_services", @@ -106,11 +109,17 @@ def test_installs_the_extension_with_the_workspace_and_its_claude_and_gpt_models { "id": "system.ai.claude-sonnet-5-5", "name": "Claude Sonnet 5.5", - "api": "chat_completions", + "family": "claude-sonnet-5-5", + "api": "anthropic", "maxInputTokens": 872_000, "maxOutputTokens": 128_000, }, - {"id": "system.ai.gpt-6-luna", "name": "GPT 6 Luna", "api": "responses"}, + { + "id": "system.ai.gpt-6-luna", + "name": "GPT 6 Luna", + "family": "gpt-6-luna", + "api": "responses", + }, ], } printed = _printed(capsys) @@ -132,14 +141,142 @@ def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): ), ) - models, _ = vscode.model_entries(WS, "token") + models = vscode.model_entries(STATE) assert [model["name"] for model in models] == [ - "system.ai.gpt-6-luna", "main.models.gpt-6-luna", + "system.ai.gpt-6-luna", "GPT 5.5", ] + def test_a_provider_services_claude_targets_keep_their_ids_and_get_a_200k_cap( + self, code, monkeypatch, capsys + ): + service = { + "name": "main.default.bedrock", + "provider_type": "amazon_bedrock", + "targets": [ + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "global.anthropic.claude-opus-4-8", + "openai.gpt-oss-120b-1:0", + ], + } + monkeypatch.setattr( + databricks_module, "list_model_provider_services", lambda *a, **k: ([service], None) + ) + + assert vscode.configure(STATE, provider="main.default.bedrock") + + assert code.config["provider"] == "main.default.bedrock" + assert code.config["models"] == [ + { + "id": "global.anthropic.claude-opus-4-8", + "name": "global.anthropic.claude-opus-4-8", + "family": "claude-opus-4-8", + "api": "anthropic", + "maxInputTokens": 168_000, + "maxOutputTokens": 32_000, + }, + { + "id": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "name": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "family": "claude-sonnet-4-5", + "api": "anthropic", + "maxInputTokens": 168_000, + "maxOutputTokens": 32_000, + }, + ] + printed = _printed(capsys) + assert "with 2 models from main.default.bedrock" in printed + assert "that VS Code can't use: openai.gpt-oss-120b-1:0" in printed + + def test_a_openai_provider_services_models_on_responses_and_chat_completions( + self, code, monkeypatch + ): + service = { + "name": "main.default.openai", + "provider_type": "openai", + "allow_all_targets": True, + "targets": [], + } + monkeypatch.setattr( + databricks_module, "list_model_provider_services", lambda *a, **k: ([service], None) + ) + monkeypatch.setattr( + copilot, + "list_provider_service_models", + lambda *args: (["gpt-5.5", "gpt-4.1"], None), + ) + + assert vscode.configure(STATE, provider="main.default.openai") + + assert code.config["models"] == [ + {"id": "gpt-5.5", "name": "gpt-5.5", "family": "gpt-5.5", "api": "responses"}, + {"id": "gpt-4.1", "name": "gpt-4.1", "family": "gpt-4.1", "api": "completions"}, + ] + + @pytest.mark.parametrize( + ("service", "error"), + [ + ({"provider_type": "gemini_enterprise"}, "which vscode can't route to"), + ({"provider_type": "anthropic", "relayed": True}, "only Claude Code can use"), + ( + {"provider_type": "openai", "allow_all_targets": True}, + "Could not list the models of 'main.default.mps': forbidden.", + ), + ], + ) + def test_keeps_the_installed_setup_for_a_provider_service_it_cant_use( + self, code, monkeypatch, capsys, service, error + ): + service = {"name": "main.default.mps", "targets": [], **service} + monkeypatch.setattr( + databricks_module, "list_model_provider_services", lambda *a, **k: ([service], None) + ) + monkeypatch.setattr( + copilot, "list_provider_service_models", lambda *args: ([], "forbidden") + ) + + assert not vscode.configure(STATE, provider="main.default.mps") + + assert code.calls == [] + printed = _printed(capsys) + assert error in printed + assert "VS Code was left unchanged." in printed + + def test_a_model_locations_services_with_a_200k_cap(self, code, monkeypatch): + listed: list[str | None] = [] + + def listing(workspace, token, *, parent_schema=None): + listed.append(parent_schema) + return AnthropicModelCatalog( + model_ids=["main.models.claude-opus-5-5", "main.models.team-assistant"], + model_id_to_display_name={}, + ) + + monkeypatch.setattr(copilot, "list_anthropic_model_catalog", listing) + + assert vscode.configure(STATE, location="main.models") + + assert listed == ["main.models"] + assert "provider" not in code.config + assert code.config["models"] == [ + { + "id": "main.models.claude-opus-5-5", + "name": "Claude Opus 5.5", + "family": "claude-opus-5-5", + "api": "anthropic", + "maxInputTokens": 168_000, + "maxOutputTokens": 32_000, + }, + { + "id": "main.models.team-assistant", + "name": "Team Assistant", + "family": "team-assistant", + "api": "anthropic", + }, + ] + def test_reruns_reinstall_with_the_new_setup(self, code): vscode.configure(STATE) code.calls.clear() @@ -322,7 +459,7 @@ def test_configure_agents_vscode_refuses_a_workspace_the_machine_isnt_managed_fo "ucode.cli.managed_state_workspace", lambda: "https://managed.cloud.databricks.com" ) monkeypatch.setattr("ucode.cli.set_current_workspace", selected.append) - monkeypatch.setattr(vscode, "configure", lambda state: pytest.fail("VS Code was set up")) + monkeypatch.setattr(vscode, "configure", lambda state, **_: pytest.fail("VS Code was set up")) result = CliRunner().invoke(app, ["configure", "--agents", "vscode", "--workspace", WS]) @@ -348,7 +485,11 @@ def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, lambda **kwargs: pytest.fail("no model agents to configure"), ) monkeypatch.setattr("ucode.cli.load_state", lambda: STATE) - monkeypatch.setattr(vscode, "configure", lambda state: configured.append(state) or True) + monkeypatch.setattr( + vscode, + "configure", + lambda state, **source: configured.append({**state, **source}) or True, + ) args = [ "configure", @@ -356,6 +497,8 @@ def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, "vscode", "--workspace", WS, + "--provider", + "main.default.mps", *(["--dry-run"] if dry_run else []), ] result = CliRunner().invoke(app, args) @@ -363,5 +506,36 @@ def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, assert result.exit_code == 0, result.output assert shared[0]["tools"] == [] # A dry run saves no workspace, so VS Code would be set up from stale or missing state. - assert configured == ([] if dry_run else [STATE]) + source = {"provider": "main.default.mps", "location": None} + assert configured == ([] if dry_run else [{**STATE, **source}]) assert ("VS Code isn't changed" in result.output) == dry_run + + +@pytest.mark.parametrize( + ("args", "error"), + [ + (["--agents", "claude", "--provider", "a.b.c"], "work only with `--agents vscode`"), + (["--provider", "a.b.c"], "work only with `--agents vscode`"), + # A model agent alongside VS Code would silently keep the workspace's own models. + ( + ["--agents", "claude,vscode", "--provider", "a.b.c"], + "work only with `--agents vscode`", + ), + ( + ["--agents", "vscode", "--provider", "a.b.c", "--model-location", "a.b"], + "cannot be used together", + ), + (["--agents", "vscode", "--model-location", "a.b.c"], "must be `.`"), + ], +) +def test_configure_refuses_a_model_source_vscode_cant_take(monkeypatch, args, error): + monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda: None) + monkeypatch.setattr( + "ucode.cli._configure_shared_workspace_states", + lambda *a, **k: pytest.fail("the workspace was configured"), + ) + + result = CliRunner().invoke(app, ["configure", "--workspace", WS, *args]) + + assert result.exit_code != 0 + assert error in " ".join(result.output.split()) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index bcf26da70..20fc922b9 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1474,6 +1474,12 @@ def test_reports_a_codex_listing_failure(self, monkeypatch): ("main.ucode.claude_haiku", "claude_haiku"), ("gpt-5.5", "gpt-5.5"), ("main.default.gpt-5.5", "gpt-5.5"), + # An Azure deployment names its region with a slash. + ("eu/gpt-5.6-luna", "gpt-5.6-luna"), + ("us/gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"), + ("mistralai/Mistral-7B", "mistral-7b"), + # Only the leading segment is a region; anything deeper is left alone. + ("publishers/google/models/gemini-4", "google/models/gemini-4"), ], ) def test_bare_model_name(model, name): diff --git a/vscode-extension/README.md b/vscode-extension/README.md index b79ed3ea4..8abee96ef 100644 --- a/vscode-extension/README.md +++ b/vscode-extension/README.md @@ -1,13 +1,18 @@ # Unity Gateway for VS Code Adds Databricks AI Gateway models to the VS Code chat model picker, for Copilot Chat in Ask and Agent -mode. The extension calls the workspace's AI Gateway directly: Chat Completions for Claude, Responses -for GPT-5 and later. +mode. The extension calls the workspace's AI Gateway directly: Anthropic Messages for Claude, Responses +for GPT-5 and later, and Chat Completions for older GPT models. `ug configure --agents vscode` packs the extension with a `ug-config.json` that holds the workspace URL, -the models (id, picker name, `chat_completions` or `responses`, token limits) and a command that prints -a Databricks access token (`ug auth-token --host ...`), then installs it. To change the models or the -workspace, run `ug configure --agents vscode` again and restart VS Code. +the models (id, picker name, family, the wire to call it on - `anthropic`, `responses` or +`completions` - and token limits) and a command that prints a Databricks access token (`ug auth-token --host ...`), then installs +it. To change the models or the workspace, run `ug configure --agents vscode` again and restart VS Code. + +The models are the workspace's own (`system.ai`) by default. With `--provider ` +they are a Model Provider Service's: the config then names the service, every request carries the +`Databricks-Model-Provider-Service` header, and OpenAI models use the gateway's `/openai/v1` route. +With `--model-location ` they are the model services in that schema. The token is kept in memory until it nears expiry, and fetched again when the gateway answers 401. Token fetches and failed requests are logged to the **Unity Gateway** output channel. diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 2ba3da005..4efd792e9 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -3,13 +3,26 @@ import { join } from "node:path"; import * as vscode from "vscode"; import { version } from "../package.json"; import { Tokens } from "./auth"; -import { chatCompletionsReplies, chatCompletionsRequest, estimateTokens, responsesReplies, responsesRequest, type Reply } from "./formats"; +import { + anthropicMessagesReplies, + anthropicMessagesRequest, + chatCompletionsReplies, + chatCompletionsRequest, + estimateTokens, + responsesReplies, + responsesRequest, + type Reply, +} from "./formats"; import { streamEvents } from "./gateway"; +// ug writes the wire name of the model it configured (WIRE_* in src/ucode/agents/copilot.py). +const APIS = ["anthropic", "completions", "responses"] as const; + export interface ModelSetting { readonly id: string; readonly name?: string; - readonly api: "chat_completions" | "responses"; + readonly family?: string; + readonly api: (typeof APIS)[number]; readonly maxInputTokens?: number; readonly maxOutputTokens?: number; } @@ -18,8 +31,13 @@ export interface Settings { readonly workspace: string; readonly models: readonly ModelSetting[]; readonly authCommand: string; + /** The Model Provider Service the models belong to, or "" for the workspace's own models. */ + readonly provider: string; } +/** Routes a request to a Model Provider Service; its models aren't served without it. */ +export const PROVIDER_HEADER = "Databricks-Model-Provider-Service"; + const tokenLimits = (model: ModelSetting): Partial => Object.fromEntries( (["maxInputTokens", "maxOutputTokens"] as const).flatMap((key) => { @@ -49,15 +67,17 @@ export const readConfig = (text: string | undefined): Settings => { models: (Array.isArray(models) ? models : []) .filter( (model): model is ModelSetting => - typeof model?.id === "string" && model.id !== "" && (model.api === "chat_completions" || model.api === "responses"), + typeof model?.id === "string" && model.id !== "" && APIS.includes(model.api), ) .map((model) => ({ id: model.id, ...(typeof model.name === "string" && model.name.trim() ? { name: model.name.trim() } : {}), + ...(typeof model.family === "string" && model.family ? { family: model.family } : {}), api: model.api, ...tokenLimits(model), })), authCommand: String(get("authCommand") ?? ""), + provider: typeof get("provider") === "string" ? String(get("provider")) : "", }; }; @@ -82,18 +102,18 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider this.log("No Databricks models: run `ug configure --agents vscode` to set them up."); return []; } - return models.map(({ id, name, api, maxInputTokens, maxOutputTokens }) => ({ + return models.map(({ id, name, family, api, maxInputTokens, maxOutputTokens }) => ({ id, api, name: name ?? id, // Copilot Chat tunes its prompts and tools by family, so this is the model's own name. - family: id.split(".").pop() ?? id, + family: family ?? id.split(".").pop() ?? id, version, detail: "Databricks", // The gateway reports no limits, so a model ug configure has none for gets Copilot's 128K default. maxInputTokens: maxInputTokens ?? 128_000, maxOutputTokens: maxOutputTokens ?? 32_000, - capabilities: { toolCalling: true, imageInput: api === "responses" || id.includes("claude") }, + capabilities: { toolCalling: true, imageInput: api !== "completions" || id.includes("claude") }, })); } @@ -104,17 +124,29 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider progress: vscode.Progress, token: vscode.CancellationToken, ): Promise { - const { workspace, authCommand } = this.settings(); + const { workspace, authCommand, provider } = this.settings(); this.tokens ??= new Tokens(authCommand, (expiresAt) => this.log(`Fetched a Databricks token (expires ${new Date(expiresAt).toISOString()})`)); const turn = { model: model.id, messages, tools: options.tools ?? [], toolRequired: options.toolMode === vscode.LanguageModelChatToolMode.Required }; const abort = new AbortController(); const cancelled = token.onCancellationRequested(() => abort.abort()); + const headers = { ...this.headers(), ...(provider ? { [PROVIDER_HEADER]: provider } : {}) }; + const tokens = this.tokens; + const stream = (path: string, body: object, extra: Record = {}) => + streamEvents(workspace, path, body, tokens, { ...headers, ...extra }, abort.signal); + // The gateway serves a Model Provider Service's OpenAI models only on its OpenAI route. + const openai = provider ? "/ai-gateway/openai/v1" : "/ai-gateway/mlflow/v1"; + const apis: Record AsyncIterable> = { + anthropic: () => + anthropicMessagesReplies( + stream("/ai-gateway/anthropic/v1/messages", anthropicMessagesRequest(turn, model.maxOutputTokens), { + "anthropic-version": "2023-06-01", + }), + ), + responses: () => responsesReplies(stream(`${openai}/responses`, responsesRequest(turn))), + completions: () => chatCompletionsReplies(stream(`${openai}/chat/completions`, chatCompletionsRequest(turn))), + }; try { - const replies: AsyncIterable = - model.api === "responses" - ? responsesReplies(streamEvents(workspace, "responses", responsesRequest(turn), this.tokens, this.headers(), abort.signal)) - : chatCompletionsReplies(streamEvents(workspace, "chat/completions", chatCompletionsRequest(turn), this.tokens, this.headers(), abort.signal)); - for await (const reply of replies) { + for await (const reply of apis[model.api]()) { progress.report(reply); } } catch (error) { diff --git a/vscode-extension/src/formats.ts b/vscode-extension/src/formats.ts index 3e2be406f..defbe7a26 100644 --- a/vscode-extension/src/formats.ts +++ b/vscode-extension/src/formats.ts @@ -27,11 +27,16 @@ const textOf = (part: unknown): string => { return ""; }; -const imageUrl = (part: unknown): string | undefined => +const imageData = (part: unknown): { mimeType: string; data: string } | undefined => part instanceof vscode.LanguageModelDataPart && part.mimeType.startsWith("image/") - ? `data:${part.mimeType};base64,${Buffer.from(part.data).toString("base64")}` + ? { mimeType: part.mimeType, data: Buffer.from(part.data).toString("base64") } : undefined; +const imageUrl = (part: unknown): string | undefined => { + const image = imageData(part); + return image && `data:${image.mimeType};base64,${image.data}`; +}; + // The gateway requires a schema on every function tool; VS Code omits it for tools that take no input. const parametersOf = (tool: vscode.LanguageModelChatTool): object => tool.inputSchema ?? { type: "object", properties: {} }; @@ -48,11 +53,12 @@ const toolCall = (callId: string, name: string, args: string): vscode.LanguageMo return new vscode.LanguageModelToolCallPart(callId, name, input); }; -/** Why the model stopped short, from a Responses `incomplete_details.reason` or a Chat Completions `finish_reason`. */ +/** Why the model stopped short, from a Responses `incomplete_details.reason`, a Chat Completions + * `finish_reason` or an Anthropic `stop_reason`. */ const stoppedEarly = (reason: string | undefined): Error => - reason === "max_output_tokens" || reason === "length" + reason === "max_output_tokens" || reason === "length" || reason === "max_tokens" ? new Error("The model stopped at its output token limit before finishing the response.") - : reason === "content_filter" + : reason === "content_filter" || reason === "refusal" ? new Error("The model's response was stopped by a content filter.") : new Error(`The model stopped before finishing the response (${reason ?? "no reason given"}).`); @@ -188,10 +194,14 @@ const functionOutput = (content: readonly unknown[]): string | object[] => { // Stateless: the gateway keeps no response state, so every turn resends the whole conversation. export function responsesRequest(turn: Turn): object { + const instructions: string[] = []; const input = turn.messages.flatMap((message): object[] => { const { results, calls, text, images } = sortParts(message.content); if (Number(message.role) === SYSTEM_ROLE) { - return [{ role: "system", content: text }]; + // Microsoft Foundry rejects a system-role input item outright, whatever its content, while + // every OpenAI-compatible surface takes the prompt here. + instructions.push(text); + return []; } if (message.role === vscode.LanguageModelChatMessageRole.Assistant) { return [ @@ -210,6 +220,7 @@ export function responsesRequest(turn: Turn): object { }); return { model: turn.model, + ...(instructions.length ? { instructions: instructions.join("\n\n") } : {}), input, stream: true, store: false, @@ -248,3 +259,109 @@ export async function* responsesReplies(events: AsyncIterable): AsyncGe } yield* calls; } + +const textBlocks = (text: string): object[] => (text ? [{ type: "text", text }] : []); + +const imageBlocks = (parts: readonly unknown[]): object[] => + parts.flatMap((part) => { + const image = imageData(part); + return image ? [{ type: "image", source: { type: "base64", media_type: image.mimeType, data: image.data } }] : []; + }); + +const CACHE_POINT = { cache_control: { type: "ephemeral" } }; + +/** Marks the last block, so the gateway caches the prompt up to it for the next turn. */ +const cacheUpTo = (blocks: object[]): void => { + if (blocks.length) { + blocks[blocks.length - 1] = { ...blocks[blocks.length - 1], ...CACHE_POINT }; + } +}; + +export function anthropicMessagesRequest(turn: Turn, maxTokens: number): object { + const system: object[] = []; + const messages: { role: "user" | "assistant"; content: object[] }[] = []; + for (const message of turn.messages) { + const { results, calls, text } = sortParts(message.content); + if (Number(message.role) === SYSTEM_ROLE) { + system.push(...textBlocks(text)); + continue; + } + const role = message.role === vscode.LanguageModelChatMessageRole.Assistant ? "assistant" : "user"; + const content = + role === "assistant" + ? [...textBlocks(text), ...calls.map((call) => ({ type: "tool_use", id: call.callId, name: call.name, input: call.input }))] + : [ + // Tool results have to come first in the user message that follows the calls. + ...results.map((result) => { + const blocks = [...textBlocks(result.content.map(textOf).join("")), ...imageBlocks(result.content)]; + return { type: "tool_result", tool_use_id: result.callId, ...(blocks.length ? { content: blocks } : {}) }; + }), + ...textBlocks(text), + ...imageBlocks(message.content), + ]; + const last = messages[messages.length - 1]; + if (!content.length) { + continue; + } else if (last?.role === role) { + last.content.push(...content); + } else { + messages.push({ role, content }); + } + } + cacheUpTo(system); + cacheUpTo(messages[messages.length - 1]?.content ?? []); + return { + model: turn.model, + max_tokens: maxTokens, + ...(system.length ? { system } : {}), + messages, + stream: true, + ...(turn.tools.length + ? { + tools: turn.tools.map((tool) => ({ name: tool.name, description: tool.description, input_schema: parametersOf(tool) })), + ...(turn.toolRequired ? { tool_choice: { type: "any" } } : {}), + } + : {}), + }; +} + +interface AnthropicEvent { + type: string; + index?: number; + content_block?: { type?: string; id?: string; name?: string }; + delta?: { type?: string; text?: string; partial_json?: string; stop_reason?: string | null }; + error?: { message?: string }; +} + +const ANTHROPIC_FINISHED = new Set(["end_turn", "tool_use", "stop_sequence"]); + +export async function* anthropicMessagesReplies(events: AsyncIterable): AsyncGenerator { + // Held back until the message stops, so a cut-off stream runs no tool. + const calls = new Map(); + let stopped = false; + for await (const raw of events) { + const event = raw as AnthropicEvent; + if (event.type === "content_block_start" && event.content_block?.type === "tool_use") { + calls.set(event.index ?? 0, { id: event.content_block.id ?? "", name: event.content_block.name ?? "", args: "" }); + } else if (event.type === "content_block_delta" && event.delta?.type === "text_delta" && event.delta.text) { + yield new vscode.LanguageModelTextPart(event.delta.text); + } else if (event.type === "content_block_delta" && event.delta?.type === "input_json_delta") { + const call = calls.get(event.index ?? 0); + if (call) { + call.args += event.delta.partial_json ?? ""; + } + } else if (event.type === "message_delta" && event.delta?.stop_reason && !ANTHROPIC_FINISHED.has(event.delta.stop_reason)) { + throw stoppedEarly(event.delta.stop_reason); + } else if (event.type === "message_stop") { + stopped = true; + } else if (event.type === "error") { + throw new Error(event.error?.message ?? "The model response failed."); + } + } + if (!stopped) { + throw cutOff(); + } + for (const call of calls.values()) { + yield toolCall(call.id, call.name, call.args); + } +} diff --git a/vscode-extension/src/gateway.ts b/vscode-extension/src/gateway.ts index 5db9b7959..d07da86bd 100644 --- a/vscode-extension/src/gateway.ts +++ b/vscode-extension/src/gateway.ts @@ -1,34 +1,44 @@ import type { Tokens } from "./auth"; -/** POSTs a streaming request to the workspace's AI Gateway and yields each server-sent event's JSON. */ +/** POSTs a streaming request to `path` on the workspace's AI Gateway and yields each server-sent event's JSON. */ export async function* streamEvents( workspace: string, - path: "chat/completions" | "responses", + path: string, body: object, tokens: Tokens, headers: Record, signal: AbortSignal, ): AsyncGenerator { - const post = (token: string) => - fetch(`${workspace}/ai-gateway/mlflow/v1/${path}`, { + const post = (token: string, payload: object) => + fetch(`${workspace}${path}`, { method: "POST", headers: { ...headers, Authorization: `Bearer ${token}`, "Content-Type": "application/json" }, - body: JSON.stringify(body), + body: JSON.stringify(payload), signal, }); let token = await tokens.get(); - let response = await post(token); + let response = await post(token, body); if (response.status === 401) { await response.body?.cancel(); token = await tokens.renew(token); - response = await post(token); + response = await post(token, body); } if (response.status === 401) { await response.body?.cancel(); throw new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry."); } if (!response.ok || !response.body) { - throw new Error(`AI Gateway returned ${response.status}: ${(await response.text()).slice(0, 300)}`); + const detail = (await response.text()).slice(0, 300); + // A model's real output limit can be below the one ug configured, and Anthropic Messages + // requires us to send one; the gateway's refusal names the limit to use instead. + const allowed = Number(/max_tokens: \d+ > (\d+)/.exec(detail)?.[1]); + if (!allowed) { + throw new Error(`AI Gateway returned ${response.status}: ${detail}`); + } + response = await post(token, { ...body, max_tokens: allowed }); + if (!response.ok || !response.body) { + throw new Error(`AI Gateway returned ${response.status}: ${(await response.text()).slice(0, 300)}`); + } } const decoder = new TextDecoder(); let buffer = ""; diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 537380ec8..af3d44225 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -4,7 +4,7 @@ import { afterEach, describe, it } from "node:test"; import type * as vscode from "vscode"; import packageJson from "../package.json"; import { Tokens } from "../src/auth"; -import { DatabricksModels, readConfig, userAgent, type Settings } from "../src/extension"; +import { DatabricksModels, PROVIDER_HEADER, readConfig, userAgent, type Settings } from "../src/extension"; import { streamEvents } from "../src/gateway"; import { collect, message, rejectsWith, text, tokenCommand, USER } from "./support"; import { LanguageModelChatToolMode, LanguageModelTextPart } from "./vscode"; @@ -13,10 +13,16 @@ const sse = (...data: string[]) => data.map((item) => `data: ${item}\n\n`).join( /** Replaces fetch with one that answers each request with the next status and SSE body. */ const stubFetch = (...answers: [number, string][]) => { - const requests: { url: string; authorization: string; userAgent: string; body: Record }[] = []; + const requests: { url: string; authorization: string; userAgent: string; provider: string | null; body: Record }[] = []; globalThis.fetch = (async (url: string, init: RequestInit) => { const headers = new Headers(init.headers); - requests.push({ url, authorization: headers.get("authorization") ?? "", userAgent: headers.get("user-agent") ?? "", body: JSON.parse(String(init.body)) }); + requests.push({ + url, + authorization: headers.get("authorization") ?? "", + userAgent: headers.get("user-agent") ?? "", + provider: headers.get(PROVIDER_HEADER), + body: JSON.parse(String(init.body)), + }); const [status, body] = answers[Math.min(requests.length - 1, answers.length - 1)] ?? [500, ""]; // Split the body mid-line so the reader has to join chunks. const bytes = new TextEncoder().encode(body); @@ -35,33 +41,53 @@ describe("streamEvents", () => { it("yields each event's JSON until [DONE]", async () => { stubFetch([200, sse('{"a":1}', '{"b":2}', "[DONE]", '{"c":3}')]); const tokens = new Tokens(tokenCommand(["t"]).command); - assert.deepEqual(await collect(streamEvents("https://ws", "responses", {}, tokens, {}, new AbortController().signal)), [{ a: 1 }, { b: 2 }]); + assert.deepEqual(await collect(streamEvents("https://ws", "/r", {}, tokens, {}, new AbortController().signal)), [{ a: 1 }, { b: 2 }]); }); it("keeps a last event that has no trailing newline", async () => { stubFetch([200, `data: {"a":1}\n\ndata: {"b":2}`]); const tokens = new Tokens(tokenCommand(["t"]).command); - assert.deepEqual(await collect(streamEvents("https://ws", "responses", {}, tokens, {}, new AbortController().signal)), [{ a: 1 }, { b: 2 }]); + assert.deepEqual(await collect(streamEvents("https://ws", "/r", {}, tokens, {}, new AbortController().signal)), [{ a: 1 }, { b: 2 }]); }); it("retries once with a new token after a 401, then asks the user to sign in", async () => { let requests = stubFetch([401, ""], [200, sse("[DONE]")]); - await collect(streamEvents("https://ws", "chat/completions", {}, new Tokens(tokenCommand(["fresh"], "rejected").command), {}, new AbortController().signal)); + await collect(streamEvents("https://ws", "/ai-gateway/mlflow/v1/chat/completions", {}, new Tokens(tokenCommand(["fresh"], "rejected").command), {}, new AbortController().signal)); assert.deepEqual(requests.map((request) => request.authorization), ["Bearer rejected", "Bearer fresh"]); assert.equal(requests[0]?.url, "https://ws/ai-gateway/mlflow/v1/chat/completions"); requests = stubFetch([401, ""]); await rejectsWith( - collect(streamEvents("https://ws", "responses", {}, new Tokens(tokenCommand(["a", "b"]).command), {}, new AbortController().signal)), + collect(streamEvents("https://ws", "/r", {}, new Tokens(tokenCommand(["a", "b"]).command), {}, new AbortController().signal)), "Run `ug configure --agents vscode` to sign in again", ); assert.equal(requests.length, 2); }); + it("retries once with the output limit the gateway allows, and reports any other error", async () => { + const tooBig = [400, '{"message":"max_tokens: 32000 > 8192, which is the maximum allowed number of output tokens for claude-3-5-sonnet"}'] as [number, string]; + let requests = stubFetch(tooBig, [200, sse('{"a":1}', "[DONE]")]); + const events = await collect( + streamEvents("https://ws", "/m", { model: "claude-3-5-sonnet", max_tokens: 32_000 }, new Tokens(tokenCommand(["t"]).command), {}, new AbortController().signal), + ); + assert.deepEqual(events, [{ a: 1 }]); + assert.deepEqual( + requests.map((request) => request.body.max_tokens), + [32_000, 8_192], + ); + + requests = stubFetch(tooBig, tooBig); + await rejectsWith( + collect(streamEvents("https://ws", "/m", { max_tokens: 32_000 }, new Tokens(tokenCommand(["t"]).command), {}, new AbortController().signal)), + "AI Gateway returned 400: {\"message\":\"max_tokens: 32000 > 8192", + ); + assert.equal(requests.length, 2); + }); + it("reports other gateway errors with their body", async () => { stubFetch([400, '{"message":"bad tools"}']); await rejectsWith( - collect(streamEvents("https://ws", "responses", {}, new Tokens(tokenCommand(["t"]).command), {}, new AbortController().signal)), + collect(streamEvents("https://ws", "/r", {}, new Tokens(tokenCommand(["t"]).command), {}, new AbortController().signal)), 'AI Gateway returned 400: {"message":"bad tools"}', ); }); @@ -82,32 +108,38 @@ describe("userAgent", () => { }); describe("readConfig", () => { - it("keeps well-formed models, names and token limits and trims the workspace URL", () => { + it("keeps well-formed models, names, families and token limits, the provider, and trims the workspace URL", () => { const config = { workspace: "https://ws/", authCommand: "ug auth-token", + provider: "main.default.mps", models: [ { id: "a", api: "responses" }, + { id: "gpt-5.5", family: "gpt-5.5", api: "responses" }, + { id: "e", api: "anthropic" }, { id: "b", api: "messages" }, { api: "responses" }, null, - { id: "c", name: " Model C ", api: "chat_completions", maxInputTokens: 936_000, maxOutputTokens: "32000", extra: true }, - { id: "d", name: 7, api: "chat_completions", maxInputTokens: 0, maxOutputTokens: 1.5 }, + { id: "c", name: " Model C ", api: "completions", maxInputTokens: 936_000, maxOutputTokens: "32000", extra: true }, + { id: "d", name: 7, api: "completions", maxInputTokens: 0, maxOutputTokens: 1.5 }, ], }; assert.deepEqual(readConfig(JSON.stringify(config)), { workspace: "https://ws", authCommand: "ug auth-token", + provider: "main.default.mps", models: [ { id: "a", api: "responses" }, - { id: "c", name: "Model C", api: "chat_completions", maxInputTokens: 936_000 }, - { id: "d", api: "chat_completions" }, + { id: "gpt-5.5", family: "gpt-5.5", api: "responses" }, + { id: "e", api: "anthropic" }, + { id: "c", name: "Model C", api: "completions", maxInputTokens: 936_000 }, + { id: "d", api: "completions" }, ], }); }); it("gives no models for a missing or malformed file", () => { - const none = { workspace: "", authCommand: "", models: [] }; + const none = { workspace: "", authCommand: "", models: [], provider: "" }; for (const text of [undefined, "", "{", "[]", "null", '"text"']) { assert.deepEqual(readConfig(text), none, String(text)); } @@ -115,12 +147,14 @@ describe("readConfig", () => { }); describe("DatabricksModels", () => { - const settings = (command: string): Settings => ({ + const settings = (command: string, provider = ""): Settings => ({ workspace: "https://ws", authCommand: command, + provider, models: [ - { id: "system.ai.claude-sonnet-5-5", name: "Claude Sonnet 5.5", api: "chat_completions", maxInputTokens: 872_000, maxOutputTokens: 128_000 }, + { id: "system.ai.claude-sonnet-5-5", name: "Claude Sonnet 5.5", api: "anthropic", maxInputTokens: 872_000, maxOutputTokens: 128_000 }, { id: "system.ai.gpt-6-luna", api: "responses" }, + { id: "gpt-4.1", family: "gpt-4.1", api: "completions" }, ], }); @@ -130,23 +164,24 @@ describe("DatabricksModels", () => { assert.deepEqual( models.map((model) => [model.id, model.name, model.family, model.api, model.maxInputTokens, model.maxOutputTokens]), [ - ["system.ai.claude-sonnet-5-5", "Claude Sonnet 5.5", "claude-sonnet-5-5", "chat_completions", 872_000, 128_000], + ["system.ai.claude-sonnet-5-5", "Claude Sonnet 5.5", "claude-sonnet-5-5", "anthropic", 872_000, 128_000], ["system.ai.gpt-6-luna", "system.ai.gpt-6-luna", "gpt-6-luna", "responses", 128_000, 32_000], + ["gpt-4.1", "gpt-4.1", "gpt-4.1", "completions", 128_000, 32_000], ], ); assert.deepEqual(await new DatabricksModels(() => ({ ...settings(""), models: [] }), (line) => logs.push(line)).provideLanguageModelChatInformation(), []); assert.deepEqual(logs, ["No Databricks models: run `ug configure --agents vscode` to set them up."]); }); - it("sends each model's turns to its API with the User-Agent", async () => { - const requests = stubFetch( - [200, sse('{"choices":[{"delta":{"content":"hello"},"finish_reason":"stop"}]}', "[DONE]")], - [200, sse('{"type":"response.output_text.delta","delta":"hi"}', '{"type":"response.completed","response":{}}')], - ); - const provider = new DatabricksModels(() => settings(tokenCommand(["t"]).command), () => {}, () => ({ "User-Agent": "ucode/0 githubcopilotchat" })); - const [claude, gpt] = await provider.provideLanguageModelChatInformation(); + const answers: [number, string][] = [ + [200, sse('{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}', '{"type":"message_stop"}')], + [200, sse('{"type":"response.output_text.delta","delta":"hi"}', '{"type":"response.completed","response":{}}')], + [200, sse('{"choices":[{"delta":{"content":"hey"},"finish_reason":"stop"}]}', "[DONE]")], + ]; + + const askEach = async (provider: DatabricksModels) => { const reported: unknown[] = []; - for (const model of [claude!, gpt!]) { + for (const model of await provider.provideLanguageModelChatInformation()) { await provider.provideLanguageModelChatResponse( model, [message(USER, text("hi"))], @@ -155,13 +190,35 @@ describe("DatabricksModels", () => { { isCancellationRequested: false, onCancellationRequested: () => ({ dispose() {} }) } as unknown as vscode.CancellationToken, ); } + return reported; + }; + + it("sends each model's turns to its API with the User-Agent and its output limit", async () => { + const requests = stubFetch(...answers); + const provider = new DatabricksModels(() => settings(tokenCommand(["t"]).command), () => {}, () => ({ "User-Agent": "ucode/0 githubcopilotchat" })); + const reported = await askEach(provider); + assert.deepEqual( + requests.map((request) => [request.url, request.body.model, request.userAgent, request.provider]), + [ + ["https://ws/ai-gateway/anthropic/v1/messages", "system.ai.claude-sonnet-5-5", "ucode/0 githubcopilotchat", null], + ["https://ws/ai-gateway/mlflow/v1/responses", "system.ai.gpt-6-luna", "ucode/0 githubcopilotchat", null], + ["https://ws/ai-gateway/mlflow/v1/chat/completions", "gpt-4.1", "ucode/0 githubcopilotchat", null], + ], + ); + assert.equal(requests[0]?.body.max_tokens, 128_000); + assert.deepEqual(reported, [new LanguageModelTextPart("hello"), new LanguageModelTextPart("hi"), new LanguageModelTextPart("hey")]); + }); + + it("routes a Model Provider Service's models with its header, and OpenAI ones on the OpenAI route", async () => { + const requests = stubFetch(...answers); + await askEach(new DatabricksModels(() => settings(tokenCommand(["t"]).command, "main.default.mps"), () => {})); assert.deepEqual( - requests.map((request) => [request.url, request.body.model, request.userAgent]), + requests.map((request) => [request.url, request.provider]), [ - ["https://ws/ai-gateway/mlflow/v1/chat/completions", "system.ai.claude-sonnet-5-5", "ucode/0 githubcopilotchat"], - ["https://ws/ai-gateway/mlflow/v1/responses", "system.ai.gpt-6-luna", "ucode/0 githubcopilotchat"], + ["https://ws/ai-gateway/anthropic/v1/messages", "main.default.mps"], + ["https://ws/ai-gateway/openai/v1/responses", "main.default.mps"], + ["https://ws/ai-gateway/openai/v1/chat/completions", "main.default.mps"], ], ); - assert.deepEqual(reported, [new LanguageModelTextPart("hello"), new LanguageModelTextPart("hi")]); }); }); diff --git a/vscode-extension/test/formats.test.ts b/vscode-extension/test/formats.test.ts index 6abf2c81d..09f20604a 100644 --- a/vscode-extension/test/formats.test.ts +++ b/vscode-extension/test/formats.test.ts @@ -1,6 +1,14 @@ import assert from "node:assert/strict"; import { describe, it } from "node:test"; -import { chatCompletionsReplies, chatCompletionsRequest, estimateTokens, responsesReplies, responsesRequest } from "../src/formats"; +import { + anthropicMessagesReplies, + anthropicMessagesRequest, + chatCompletionsReplies, + chatCompletionsRequest, + estimateTokens, + responsesReplies, + responsesRequest, +} from "../src/formats"; import { ASSISTANT, call, collect, events, image, message, rejectsWith, result, SYSTEM, text, USER, weatherTool } from "./support"; import { LanguageModelDataPart, LanguageModelTextPart, LanguageModelToolCallPart } from "./vscode"; @@ -107,10 +115,10 @@ describe("Responses", () => { it("sends the whole conversation statelessly", () => { assert.deepEqual(responsesRequest(turn), { model: "m", + instructions: "Be brief.", stream: true, store: false, input: [ - { role: "system", content: "Be brief." }, { role: "user", content: [{ type: "input_text", text: "Weather?" }, { type: "input_image", image_url: "data:image/png;base64,Bw==", detail: "auto" }] }, { role: "assistant", content: "Checking." }, { type: "function_call", call_id: "c1", name: "get_weather", arguments: '{"city":"Paris"}' }, @@ -178,3 +186,128 @@ describe("Responses", () => { await rejectsWith(collect(responsesReplies(events({ type: "response.output_text.delta", delta: "Half an ans" }))), "ended before the model finished"); }); }); + +describe("Anthropic Messages", () => { + const cached = { cache_control: { type: "ephemeral" } }; + + it("sends the system prompt, conversation, images and tools, with cache points after the system prompt and the last message", () => { + assert.deepEqual(anthropicMessagesRequest(turn, 32_000), { + model: "m", + max_tokens: 32_000, + stream: true, + system: [{ type: "text", text: "Be brief.", ...cached }], + messages: [ + { + role: "user", + content: [ + { type: "text", text: "Weather?" }, + { type: "image", source: { type: "base64", media_type: "image/png", data: "Bw==" } }, + ], + }, + { + role: "assistant", + content: [ + { type: "text", text: "Checking." }, + { type: "tool_use", id: "c1", name: "get_weather", input: { city: "Paris" } }, + ], + }, + { role: "user", content: [{ type: "tool_result", tool_use_id: "c1", content: [{ type: "text", text: "Sunny" }], ...cached }] }, + ], + tools: [{ name: "get_weather", description: "Weather for a city", input_schema: weatherTool.inputSchema }], + tool_choice: { type: "any" }, + }); + }); + + it("puts a tool's images in its result, results before text, and joins messages of the same role", () => { + const request = anthropicMessagesRequest( + { + ...turn, + messages: [ + message(ASSISTANT, call("c1", "screenshot", {}), call("c2", "list", {})), + message(USER, text("Then say."), result("c1", text("Screenshot:"), image(9)), result("c2")), + message(USER, text("Briefly.")), + ], + }, + 100, + ); + assert.deepEqual((request as { messages: unknown[] }).messages, [ + { + role: "assistant", + content: [ + { type: "tool_use", id: "c1", name: "screenshot", input: {} }, + { type: "tool_use", id: "c2", name: "list", input: {} }, + ], + }, + { + role: "user", + content: [ + { + type: "tool_result", + tool_use_id: "c1", + content: [ + { type: "text", text: "Screenshot:" }, + { type: "image", source: { type: "base64", media_type: "image/png", data: "CQ==" } }, + ], + }, + { type: "tool_result", tool_use_id: "c2" }, + { type: "text", text: "Then say." }, + { type: "text", text: "Briefly.", ...cached }, + ], + }, + ]); + }); + + it("leaves out empty messages, the system prompt and tools when there are none", () => { + const request = anthropicMessagesRequest({ ...turn, messages: [message(USER, text("Hi")), message(ASSISTANT, text(""))], tools: [] }, 100); + assert.deepEqual(request, { model: "m", max_tokens: 100, stream: true, messages: [{ role: "user", content: [{ type: "text", text: "Hi", ...cached }] }] }); + }); + + it("streams text, drops thinking, and assembles tool calls from their fragments once the message stops", async () => { + const replies = await collect( + anthropicMessagesReplies( + events( + { type: "message_start", message: {} }, + { type: "content_block_start", index: 0, content_block: { type: "thinking", thinking: "" } }, + { type: "content_block_delta", index: 0, delta: { type: "thinking_delta", thinking: "The user wants the weather." } }, + { type: "content_block_start", index: 1, content_block: { type: "text", text: "" } }, + { type: "content_block_delta", index: 1, delta: { type: "text_delta", text: "Checking." } }, + { type: "content_block_start", index: 2, content_block: { type: "tool_use", id: "c1", name: "get_weather", input: {} } }, + { type: "content_block_delta", index: 2, delta: { type: "input_json_delta", partial_json: '{"ci' } }, + { type: "content_block_delta", index: 2, delta: { type: "input_json_delta", partial_json: 'ty":"Paris"}' } }, + { type: "content_block_start", index: 3, content_block: { type: "tool_use", id: "c2", name: "list", input: {} } }, + { type: "ping" }, + { type: "message_delta", delta: { stop_reason: "tool_use" } }, + { type: "message_stop" }, + ), + ), + ); + assert.deepEqual(replies, [ + new LanguageModelTextPart("Checking."), + new LanguageModelToolCallPart("c1", "get_weather", { city: "Paris" }), + new LanguageModelToolCallPart("c2", "list", {}), + ]); + }); + + it("fails on a token-limit stop, a refusal, any other stop reason, an error event, and a stream without message_stop", async () => { + const stop = (reason: string) => events({ type: "message_delta", delta: { stop_reason: reason } }, { type: "message_stop" }); + await rejectsWith(collect(anthropicMessagesReplies(stop("max_tokens"))), "output token limit"); + await rejectsWith(collect(anthropicMessagesReplies(stop("refusal"))), "stopped by a content filter"); + await rejectsWith(collect(anthropicMessagesReplies(stop("pause_turn"))), "(pause_turn)"); + await rejectsWith(collect(anthropicMessagesReplies(events({ type: "error", error: { type: "overloaded_error", message: "Overloaded" } }))), "Overloaded"); + const seen: unknown[] = []; + await rejectsWith( + (async () => { + for await (const reply of anthropicMessagesReplies( + events( + { type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "Half" } }, + { type: "content_block_start", index: 1, content_block: { type: "tool_use", id: "c1", name: "list", input: {} } }, + ), + )) { + seen.push(reply); + } + })(), + "ended before the model finished", + ); + assert.deepEqual(seen, [new LanguageModelTextPart("Half")]); + }); +});