From f9e87844507b2f8d37236488d68dba3b12244b70 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 18:33:00 +0000 Subject: [PATCH 1/7] Give VS Code a Model Provider Service's or a model location's models, and Claude on Anthropic Messages ug configure --agents vscode takes --provider or --model-location , picking the models the way ug copilot does. The extension sends a provider's header, uses the gateway's OpenAI route for its OpenAI models, and calls every Claude model through Anthropic Messages with prompt-cache points. Co-authored-by: Isaac --- src/ucode/agents/copilot.py | 50 ++++-- src/ucode/agents/vscode.py | 117 +++++++++---- src/ucode/cli.py | 27 ++- src/ucode/databricks.py | 4 +- src/ucode/vscode/extension/out/extension.js | 4 +- tests/test_agent_vscode.py | 185 +++++++++++++++++++- vscode-extension/README.md | 15 +- vscode-extension/src/extension.ts | 51 ++++-- vscode-extension/src/formats.ts | 122 ++++++++++++- vscode-extension/src/gateway.ts | 6 +- vscode-extension/test/extension.test.ts | 87 ++++++--- vscode-extension/test/formats.test.ts | 135 +++++++++++++- 12 files changed, 689 insertions(+), 114 deletions(-) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 614f6c7a4..35453cff0 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, @@ -279,6 +279,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 +358,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 = "ug 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 +403,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 +670,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..82a9648b0 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -5,12 +5,16 @@ 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 import json import os +import re import shutil import subprocess import tempfile @@ -21,12 +25,17 @@ from xml.sax.saxutils import escape from ucode.agents.copilot import ( - copilot_exclusion_reason, - copilot_model_limits, - uses_responses_api, + WIRE_ANTHROPIC, + WIRE_COMPLETIONS, + WIRE_RESPONSES, + CopilotModel, + build_copilot_catalog, + catalog_model_limits, + 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 +from ucode.databricks import bare_model_name, build_auth_shell_command, discover_model_services from ucode.managed_files import OS, current_os from ucode.os_compatibility import subprocess_cross_os from ucode.os_compatibility.system_cross_os import windows_program_files_dir @@ -40,6 +49,14 @@ # Each `code` command and the folder in the home directory it installs extensions under. CODE_COMMANDS = (("code", ".vscode"), ("code-insiders", ".vscode-insiders")) _CODE_TIMEOUT_S = 120 +# The extension's `api` for each wire API of Copilot's catalog. +_APIS = { + WIRE_ANTHROPIC: "anthropic_messages", + WIRE_RESPONSES: "responses", + WIRE_COMPLETIONS: "chat_completions", +} +# Anthropic and Bedrock Claude ids end in a release date and, on Bedrock, a version. +_RELEASE_SUFFIX = re.compile(r"(?:-\d{8})?(?:-v\d+(?::\d+)?)?$") def _install_locations(command: str) -> list[Path]: @@ -161,36 +178,60 @@ 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 _family(model: str) -> str: + """The model's own name, e.g. ``us.anthropic.claude-sonnet-4-5-20250929-v1:0`` -> + ``claude-sonnet-4-5``.""" + name = bare_model_name(model) + return _RELEASE_SUFFIX.sub("", name, count=1) if name.startswith("claude-") else name + + +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": _family(model.id), + "api": _APIS[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) + labels = {model: model_picker_label(_family(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(provider or location) + ) for model, label in labels.items() - ], reason + ] def auth_command(state: dict) -> str: @@ -218,9 +259,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 +270,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..f2779ea1a 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,15 @@ 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: + if "vscode" not in [a.strip().lower() for a in (agents or "").split(",")]: + raise RuntimeError( + "--provider and --model-location work only with --agents vscode." + ) + 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 +4747,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..4f00068c0 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 diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index c473cfd19..51acd930d 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"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 oe=Object.create;var b=Object.defineProperty;var se=Object.getOwnPropertyDescriptor;var ae=Object.getOwnPropertyNames;var re=Object.getPrototypeOf,ie=Object.prototype.hasOwnProperty;var ce=(e,n)=>{for(var o in n)b(e,o,{get:n[o],enumerable:!0})},L=(e,n,o,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let t of ae(n))!ie.call(e,t)&&t!==o&&b(e,t,{get:()=>n[t],enumerable:!(s=se(n,t))||s.enumerable});return e};var O=(e,n,o)=>(o=e!=null?oe(re(e)):{},L(n||!e||!e.__esModule?b(o,"default",{value:e,enumerable:!0}):o,e)),le=e=>L(b({},"__esModule",{value:!0}),e);var Ce={};ce(Ce,{CONFIG_FILE:()=>Z,DatabricksModels:()=>k,PROVIDER_HEADER:()=>Q,activate:()=>Me,readConfig:()=>ee,userAgent:()=>te});module.exports=le(Ce);var V=require("node:fs"),X=require("node:path"),f=O(require("vscode"));var C="0.1.0";var I=require("node:child_process"),ge=5*6e4,pe=60*6e4,ue=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())})}),me=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()+pe},_=class{constructor(n,o=()=>{}){this.command=n;this.onFetch=o}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-ge>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??=ue(n).then(o=>{if(!o)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:o,expiresAt:me(o)},this.onFetch(this.cached.expiresAt),o}).finally(()=>{this.pending=void 0}),this.pending}};var c=O(require("vscode")),T=3,fe=1e3,v=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",$=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,q=e=>{let n=$(e);return n&&`data:${n.mimeType};base64,${n.data}`},R=e=>e.inputSchema??{type:"object",properties:{}},S=(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 c.LanguageModelToolCallPart(e,n,s)},P=e=>e==="max_output_tokens"||e==="length"||e==="max_tokens"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"||e==="refusal"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),A=()=>new Error("The response ended before the model finished. Try again."),J=(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 c.LanguageModelToolResultPart?s.content:[s];for(let r of t)n+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:v(r).length,o+=q(r)===void 0?0:1}return Math.ceil(n/4)+o*fe},E=e=>e.map(q).filter(n=>n!==void 0),j=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(v).join(""),images:E(e)});function F(e){let n=e.messages.flatMap(o=>{let{results:s,calls:t,text:r,images:d}=j(o.content);if(Number(o.role)===T)return[{role:"system",content:r}];if(o.role===c.LanguageModelChatMessageRole.Assistant){let a=t.map(l=>({id:l.callId,type:"function",function:{name:l.name,arguments:JSON.stringify(l.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>E(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(v).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...J(e,o=>({type:"function",function:{name:o.name,description:o.description,parameters:R(o)}}))}}var he=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),ye=new Set(["stop","tool_calls","function_call"]);async function*H(e){let n=[],o;for await(let s of e){let{error:t}=s;if(t)throw new Error(typeof t=="string"?t:t.message??JSON.stringify(t));for(let r of s.choices??[]){let d=he(r.delta?.content);d&&(yield new c.LanguageModelTextPart(d));for(let i of r.delta?.tool_calls??[]){let a=n[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(o=r.finish_reason,!ye.has(o)))throw P(o)}}if(!o)throw A();for(let s of n.filter(Boolean))yield S(s.id,s.name,s.args)}var ve=e=>{let n=e.map(v).join(""),o=E(e);return o.length?[...n?[{type:"input_text",text:n}]:[],...o.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function U(e){let n=e.messages.flatMap(o=>{let{results:s,calls:t,text:r,images:d}=j(o.content);if(Number(o.role)===T)return[{role:"system",content:r}];if(o.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...t.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...d.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:ve(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...J(e,o=>({type:"function",name:o.name,description:o.description,parameters:R(o),strict:!1}))}}async function*B(e){let n=[],o=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new c.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")n.push(S(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 P(t.response?.incomplete_details?.reason);if(t.type==="response.failed"||t.type==="error")throw new Error(t.response?.error?.message??t.message??"The model response failed.")}}if(!o)throw A();yield*n}var x=e=>e?[{type:"text",text:e}]:[],N=e=>e.flatMap(n=>{let o=$(n);return o?[{type:"image",source:{type:"base64",media_type:o.mimeType,data:o.data}}]:[]}),we={cache_control:{type:"ephemeral"}},D=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...we})};function W(e,n){let o=[],s=[];for(let t of e.messages){let{results:r,calls:d,text:i}=j(t.content);if(Number(t.role)===T){o.push(...x(i));continue}let a=t.role===c.LanguageModelChatMessageRole.Assistant?"assistant":"user",l=a==="assistant"?[...x(i),...d.map(g=>({type:"tool_use",id:g.callId,name:g.name,input:g.input}))]:[...r.map(g=>{let m=[...x(g.content.map(v).join("")),...N(g.content)];return{type:"tool_result",tool_use_id:g.callId,...m.length?{content:m}:{}}}),...x(i),...N(t.content)],p=s[s.length-1];if(l.length)p?.role===a?p.content.push(...l):s.push({role:a,content:l});else continue}return D(o),D(s[s.length-1]?.content??[]),{model:e.model,max_tokens:n,...o.length?{system:o}:{},messages:s,stream:!0,...e.tools.length?{tools:e.tools.map(t=>({name:t.name,description:t.description,input_schema:R(t)})),...e.toolRequired?{tool_choice:{type:"any"}}:{}}:{}}}var be=new Set(["end_turn","tool_use","stop_sequence"]);async function*K(e){let n=new Map,o=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")n.set(t.index??0,{id:t.content_block.id??"",name:t.content_block.name??"",args:""});else if(t.type==="content_block_delta"&&t.delta?.type==="text_delta"&&t.delta.text)yield new c.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&&!be.has(t.delta.stop_reason))throw P(t.delta.stop_reason);if(t.type==="message_stop")o=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!o)throw A();for(let s of n.values())yield S(s.id,s.name,s.args)}async function*z(e,n,o,s,t,r){let d=m=>fetch(`${e}${n}`,{method:"POST",headers:{...t,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(o),signal:r}),i=await s.get(),a=await d(i);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await d(i)),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)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let l=new TextDecoder,p="";for await(let m of a.body){p+=l.decode(m,{stream:!0});let w=p.split(` +`);p=w.pop()??"";for(let y of w){let h=Y(y);if(h==="[DONE]")return;h&&(yield JSON.parse(h))}}let g=Y(p+l.decode());g&&g!=="[DONE]"&&(yield JSON.parse(g))}var Y=e=>e.startsWith("data:")?e.slice(5).trim():"";var _e=["anthropic_messages","chat_completions","responses"],Q="Databricks-Model-Provider-Service",xe=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let o=e[n];return typeof o=="number"&&Number.isInteger(o)&&o>0?[[n,o]]:[]})),Z="ug-config.json",ee=e=>{let n={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(n=t)}catch{}let o=t=>n[t],s=o("models");return{workspace:String(o("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&_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,...xe(t)})),authCommand:String(o("authCommand")??""),provider:typeof o("provider")=="string"?String(o("provider")):""}},te=e=>`ucode/${C} 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:i,maxInputTokens:a,maxOutputTokens:l})=>({id:t,api:i,name:r??t,family:d??t.split(".").pop()??t,version:C,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:l??32e3,capabilities:{toolCalling:!0,imageInput:i!=="chat_completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(n,o,s,t,r){let{workspace:d,authCommand:i,provider:a}=this.settings();this.tokens??=new _(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let l={model:n.id,messages:o,tools:s.tools??[],toolRequired:s.toolMode===f.LanguageModelChatToolMode.Required},p=new AbortController,g=r.onCancellationRequested(()=>p.abort()),m={...this.headers(),...a?{[Q]:a}:{}},w=this.tokens,y=(u,M,ne={})=>z(d,u,M,w,{...m,...ne},p.signal),h=a?"/ai-gateway/openai/v1":"/ai-gateway/mlflow/v1";try{let u=n.api==="anthropic_messages"?K(y("/ai-gateway/anthropic/v1/messages",W(l,n.maxOutputTokens),{"anthropic-version":"2023-06-01"})):n.api==="responses"?B(y(`${h}/responses`,U(l))):H(y(`${h}/chat/completions`,F(l)));for await(let M of u)t.report(M)}catch(u){if(!r.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,o){return G(o)}},ke=e=>{try{return(0,V.readFileSync)(e,"utf8")}catch{return}};function Me(e){let n=f.window.createOutputChannel("Unity Gateway",{log:!0}),o=ee(ke((0,X.join)(e.extensionPath,Z))),s=()=>({"User-Agent":te(f.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>o,r=>n.info(r),s);e.subscriptions.push(n,f.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,activate,readConfig,userAgent}); diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 48c55402a..3b0879694 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_messages", "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_with_a_200k_cap_and_its_name( + 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": "Claude Opus 4.8", + "family": "claude-opus-4-8", + "api": "anthropic_messages", + "maxInputTokens": 168_000, + "maxOutputTokens": 32_000, + }, + { + "id": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "name": "Claude Sonnet 4.5", + "family": "claude-sonnet-4-5", + "api": "anthropic_messages", + "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": "chat_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_messages", + "maxInputTokens": 168_000, + "maxOutputTokens": 32_000, + }, + { + "id": "main.models.team-assistant", + "name": "Team Assistant", + "family": "team-assistant", + "api": "anthropic_messages", + }, + ] + 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,31 @@ 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"), + ( + ["--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/vscode-extension/README.md b/vscode-extension/README.md index b79ed3ea4..7ee043ca1 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, `anthropic_messages`, `responses` or `chat_completions`, 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..4f556e04d 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -3,13 +3,25 @@ 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"; +const APIS = ["anthropic_messages", "chat_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 +30,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 +66,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 +101,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 !== "chat_completions" || id.includes("claude") }, })); } @@ -104,16 +123,28 @@ 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"; 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)); + model.api === "anthropic_messages" + ? anthropicMessagesReplies( + stream("/ai-gateway/anthropic/v1/messages", anthropicMessagesRequest(turn, model.maxOutputTokens), { + "anthropic-version": "2023-06-01", + }), + ) + : model.api === "responses" + ? responsesReplies(stream(`${openai}/responses`, responsesRequest(turn))) + : chatCompletionsReplies(stream(`${openai}/chat/completions`, chatCompletionsRequest(turn))); for await (const reply of replies) { progress.report(reply); } diff --git a/vscode-extension/src/formats.ts b/vscode-extension/src/formats.ts index 3e2be406f..b3fce3921 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"}).`); @@ -248,3 +254,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..4dba8f7a6 100644 --- a/vscode-extension/src/gateway.ts +++ b/vscode-extension/src/gateway.ts @@ -1,16 +1,16 @@ 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}`, { + fetch(`${workspace}${path}`, { method: "POST", headers: { ...headers, Authorization: `Bearer ${token}`, "Content-Type": "application/json" }, body: JSON.stringify(body), diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 537380ec8..8f48baecc 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,24 +41,24 @@ 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); @@ -61,7 +67,7 @@ describe("streamEvents", () => { 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,12 +88,15 @@ 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_messages" }, { id: "b", api: "messages" }, { api: "responses" }, null, @@ -98,8 +107,11 @@ describe("readConfig", () => { assert.deepEqual(readConfig(JSON.stringify(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_messages" }, { id: "c", name: "Model C", api: "chat_completions", maxInputTokens: 936_000 }, { id: "d", api: "chat_completions" }, ], @@ -107,7 +119,7 @@ describe("readConfig", () => { }); 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 +127,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_messages", maxInputTokens: 872_000, maxOutputTokens: 128_000 }, { id: "system.ai.gpt-6-luna", api: "responses" }, + { id: "gpt-4.1", family: "gpt-4.1", api: "chat_completions" }, ], }); @@ -130,23 +144,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_messages", 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", "chat_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 +170,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..067c05670 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"; @@ -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")]); + }); +}); From 1243bace0e2a34853c19d577ac521f206a498a47 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 19:43:02 +0000 Subject: [PATCH 2/7] Address review: share the Claude-name helper, reject mixed agent lists, retry a refused output limit A model's real output limit can be below the one ug configures, which Anthropic Messages rejects; the gateway names the limit it allows, so retry once with it. --provider and --model-location now refuse any other model agent in the same run, and the per-model api picks its route through a lookup. Co-authored-by: Isaac --- src/ucode/agents/copilot.py | 14 +++++++++++-- src/ucode/agents/vscode.py | 17 ++++------------ src/ucode/cli.py | 10 ++++++++-- src/ucode/vscode/extension/out/extension.js | 4 ++-- tests/test_agent_vscode.py | 9 +++++++-- vscode-extension/src/extension.ts | 22 ++++++++++----------- vscode-extension/src/gateway.ts | 20 ++++++++++++++----- vscode-extension/test/extension.test.ts | 20 +++++++++++++++++++ 8 files changed, 79 insertions(+), 37 deletions(-) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 35453cff0..8f96c55dd 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -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. diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 82a9648b0..0018590ce 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -14,7 +14,6 @@ import json import os -import re import shutil import subprocess import tempfile @@ -31,11 +30,12 @@ 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 bare_model_name, build_auth_shell_command, discover_model_services +from ucode.databricks import build_auth_shell_command, discover_model_services from ucode.managed_files import OS, current_os from ucode.os_compatibility import subprocess_cross_os from ucode.os_compatibility.system_cross_os import windows_program_files_dir @@ -55,8 +55,6 @@ WIRE_RESPONSES: "responses", WIRE_COMPLETIONS: "chat_completions", } -# Anthropic and Bedrock Claude ids end in a release date and, on Bedrock, a version. -_RELEASE_SUFFIX = re.compile(r"(?:-\d{8})?(?:-v\d+(?::\d+)?)?$") def _install_locations(command: str) -> list[Path]: @@ -178,18 +176,11 @@ def install_extension(code: Path, config: dict[str, object]) -> str: return f"installed {EXTENSION_ID} {version}" -def _family(model: str) -> str: - """The model's own name, e.g. ``us.anthropic.claude-sonnet-4-5-20250929-v1:0`` -> - ``claude-sonnet-4-5``.""" - name = bare_model_name(model) - return _RELEASE_SUFFIX.sub("", name, count=1) if name.startswith("claude-") else name - - def _model_entry(model: CopilotModel, name: str, *, external: bool) -> dict[str, object]: entry: dict[str, object] = { "id": model.id, "name": name, - "family": _family(model.id), + "family": claude_model_name(model.id), "api": _APIS[model.wire], } # Models without known limits get the extension's 128K-in, 32K-out default. @@ -223,7 +214,7 @@ def model_entries( ) -> list[dict[str, object]]: """The models of the chosen source, each with the API the extension calls it on.""" catalog = _catalog(state, provider, location) - labels = {model: model_picker_label(_family(model.id)) 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 [ diff --git a/src/ucode/cli.py b/src/ucode/cli.py index f2779ea1a..dc9479f3a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -4644,9 +4644,15 @@ def configure( 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: - if "vscode" not in [a.strip().lower() for a in (agents or "").split(",")]: + 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." + "--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.") diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index 51acd930d..92aa5d55f 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var oe=Object.create;var b=Object.defineProperty;var se=Object.getOwnPropertyDescriptor;var ae=Object.getOwnPropertyNames;var re=Object.getPrototypeOf,ie=Object.prototype.hasOwnProperty;var ce=(e,n)=>{for(var o in n)b(e,o,{get:n[o],enumerable:!0})},L=(e,n,o,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let t of ae(n))!ie.call(e,t)&&t!==o&&b(e,t,{get:()=>n[t],enumerable:!(s=se(n,t))||s.enumerable});return e};var O=(e,n,o)=>(o=e!=null?oe(re(e)):{},L(n||!e||!e.__esModule?b(o,"default",{value:e,enumerable:!0}):o,e)),le=e=>L(b({},"__esModule",{value:!0}),e);var Ce={};ce(Ce,{CONFIG_FILE:()=>Z,DatabricksModels:()=>k,PROVIDER_HEADER:()=>Q,activate:()=>Me,readConfig:()=>ee,userAgent:()=>te});module.exports=le(Ce);var V=require("node:fs"),X=require("node:path"),f=O(require("vscode"));var C="0.1.0";var I=require("node:child_process"),ge=5*6e4,pe=60*6e4,ue=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())})}),me=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()+pe},_=class{constructor(n,o=()=>{}){this.command=n;this.onFetch=o}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-ge>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??=ue(n).then(o=>{if(!o)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:o,expiresAt:me(o)},this.onFetch(this.cached.expiresAt),o}).finally(()=>{this.pending=void 0}),this.pending}};var c=O(require("vscode")),T=3,fe=1e3,v=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",$=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,q=e=>{let n=$(e);return n&&`data:${n.mimeType};base64,${n.data}`},R=e=>e.inputSchema??{type:"object",properties:{}},S=(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 c.LanguageModelToolCallPart(e,n,s)},P=e=>e==="max_output_tokens"||e==="length"||e==="max_tokens"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"||e==="refusal"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),A=()=>new Error("The response ended before the model finished. Try again."),J=(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 c.LanguageModelToolResultPart?s.content:[s];for(let r of t)n+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:v(r).length,o+=q(r)===void 0?0:1}return Math.ceil(n/4)+o*fe},E=e=>e.map(q).filter(n=>n!==void 0),j=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(v).join(""),images:E(e)});function F(e){let n=e.messages.flatMap(o=>{let{results:s,calls:t,text:r,images:d}=j(o.content);if(Number(o.role)===T)return[{role:"system",content:r}];if(o.role===c.LanguageModelChatMessageRole.Assistant){let a=t.map(l=>({id:l.callId,type:"function",function:{name:l.name,arguments:JSON.stringify(l.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>E(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(v).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...J(e,o=>({type:"function",function:{name:o.name,description:o.description,parameters:R(o)}}))}}var he=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),ye=new Set(["stop","tool_calls","function_call"]);async function*H(e){let n=[],o;for await(let s of e){let{error:t}=s;if(t)throw new Error(typeof t=="string"?t:t.message??JSON.stringify(t));for(let r of s.choices??[]){let d=he(r.delta?.content);d&&(yield new c.LanguageModelTextPart(d));for(let i of r.delta?.tool_calls??[]){let a=n[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(o=r.finish_reason,!ye.has(o)))throw P(o)}}if(!o)throw A();for(let s of n.filter(Boolean))yield S(s.id,s.name,s.args)}var ve=e=>{let n=e.map(v).join(""),o=E(e);return o.length?[...n?[{type:"input_text",text:n}]:[],...o.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function U(e){let n=e.messages.flatMap(o=>{let{results:s,calls:t,text:r,images:d}=j(o.content);if(Number(o.role)===T)return[{role:"system",content:r}];if(o.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...t.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...d.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:ve(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...J(e,o=>({type:"function",name:o.name,description:o.description,parameters:R(o),strict:!1}))}}async function*B(e){let n=[],o=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new c.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")n.push(S(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 P(t.response?.incomplete_details?.reason);if(t.type==="response.failed"||t.type==="error")throw new Error(t.response?.error?.message??t.message??"The model response failed.")}}if(!o)throw A();yield*n}var x=e=>e?[{type:"text",text:e}]:[],N=e=>e.flatMap(n=>{let o=$(n);return o?[{type:"image",source:{type:"base64",media_type:o.mimeType,data:o.data}}]:[]}),we={cache_control:{type:"ephemeral"}},D=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...we})};function W(e,n){let o=[],s=[];for(let t of e.messages){let{results:r,calls:d,text:i}=j(t.content);if(Number(t.role)===T){o.push(...x(i));continue}let a=t.role===c.LanguageModelChatMessageRole.Assistant?"assistant":"user",l=a==="assistant"?[...x(i),...d.map(g=>({type:"tool_use",id:g.callId,name:g.name,input:g.input}))]:[...r.map(g=>{let m=[...x(g.content.map(v).join("")),...N(g.content)];return{type:"tool_result",tool_use_id:g.callId,...m.length?{content:m}:{}}}),...x(i),...N(t.content)],p=s[s.length-1];if(l.length)p?.role===a?p.content.push(...l):s.push({role:a,content:l});else continue}return D(o),D(s[s.length-1]?.content??[]),{model:e.model,max_tokens:n,...o.length?{system:o}:{},messages:s,stream:!0,...e.tools.length?{tools:e.tools.map(t=>({name:t.name,description:t.description,input_schema:R(t)})),...e.toolRequired?{tool_choice:{type:"any"}}:{}}:{}}}var be=new Set(["end_turn","tool_use","stop_sequence"]);async function*K(e){let n=new Map,o=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")n.set(t.index??0,{id:t.content_block.id??"",name:t.content_block.name??"",args:""});else if(t.type==="content_block_delta"&&t.delta?.type==="text_delta"&&t.delta.text)yield new c.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&&!be.has(t.delta.stop_reason))throw P(t.delta.stop_reason);if(t.type==="message_stop")o=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!o)throw A();for(let s of n.values())yield S(s.id,s.name,s.args)}async function*z(e,n,o,s,t,r){let d=m=>fetch(`${e}${n}`,{method:"POST",headers:{...t,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(o),signal:r}),i=await s.get(),a=await d(i);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await d(i)),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)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let l=new TextDecoder,p="";for await(let m of a.body){p+=l.decode(m,{stream:!0});let w=p.split(` -`);p=w.pop()??"";for(let y of w){let h=Y(y);if(h==="[DONE]")return;h&&(yield JSON.parse(h))}}let g=Y(p+l.decode());g&&g!=="[DONE]"&&(yield JSON.parse(g))}var Y=e=>e.startsWith("data:")?e.slice(5).trim():"";var _e=["anthropic_messages","chat_completions","responses"],Q="Databricks-Model-Provider-Service",xe=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let o=e[n];return typeof o=="number"&&Number.isInteger(o)&&o>0?[[n,o]]:[]})),Z="ug-config.json",ee=e=>{let n={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(n=t)}catch{}let o=t=>n[t],s=o("models");return{workspace:String(o("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&_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,...xe(t)})),authCommand:String(o("authCommand")??""),provider:typeof o("provider")=="string"?String(o("provider")):""}},te=e=>`ucode/${C} 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:i,maxInputTokens:a,maxOutputTokens:l})=>({id:t,api:i,name:r??t,family:d??t.split(".").pop()??t,version:C,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:l??32e3,capabilities:{toolCalling:!0,imageInput:i!=="chat_completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(n,o,s,t,r){let{workspace:d,authCommand:i,provider:a}=this.settings();this.tokens??=new _(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let l={model:n.id,messages:o,tools:s.tools??[],toolRequired:s.toolMode===f.LanguageModelChatToolMode.Required},p=new AbortController,g=r.onCancellationRequested(()=>p.abort()),m={...this.headers(),...a?{[Q]:a}:{}},w=this.tokens,y=(u,M,ne={})=>z(d,u,M,w,{...m,...ne},p.signal),h=a?"/ai-gateway/openai/v1":"/ai-gateway/mlflow/v1";try{let u=n.api==="anthropic_messages"?K(y("/ai-gateway/anthropic/v1/messages",W(l,n.maxOutputTokens),{"anthropic-version":"2023-06-01"})):n.api==="responses"?B(y(`${h}/responses`,U(l))):H(y(`${h}/chat/completions`,F(l)));for await(let M of u)t.report(M)}catch(u){if(!r.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,o){return G(o)}},ke=e=>{try{return(0,V.readFileSync)(e,"utf8")}catch{return}};function Me(e){let n=f.window.createOutputChannel("Unity Gateway",{log:!0}),o=ee(ke((0,X.join)(e.extensionPath,Z))),s=()=>({"User-Agent":te(f.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>o,r=>n.info(r),s);e.subscriptions.push(n,f.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,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,o)=>{for(var n in o)b(e,n,{get:o[n],enumerable:!0})},j=(e,o,n,s)=>{if(o&&typeof o=="object"||typeof o=="function")for(let t of re(o))!ce.call(e,t)&&t!==n&&b(e,t,{get:()=>o[t],enumerable:!(s=ae(o,t))||s.enumerable});return e};var L=(e,o,n)=>(n=e!=null?se(ie(e)):{},j(o||!e||!e.__esModule?b(n,"default",{value:e,enumerable:!0}):n,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((o,n)=>{(0,I.exec)(e,{timeout:6e4,windowsHide:!0},(s,t,r)=>{s?n(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):o(String(t).trim())})}),fe=e=>{try{let{exp:o}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof o=="number")return o*1e3}catch{}return Date.now()+ue},_=class{constructor(o,n=()=>{}){this.command=o;this.onFetch=n}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(o){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==o?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(o){return this.pending??=me(o).then(n=>{if(!n)throw new Error(`Couldn't get a Databricks token: \`${o}\` printed nothing.`);return this.cached={value:n,expiresAt:fe(n)},this.onFetch(this.cached.expiresAt),n}).finally(()=>{this.pending=void 0}),this.pending}};var c=L(require("vscode")),C=3,he=1e3,w=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",D=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,$=e=>{let o=D(e);return o&&`data:${o.mimeType};base64,${o.data}`},T=e=>e.inputSchema??{type:"object",properties:{}},R=(e,o,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 ${o} with arguments that aren't a JSON object: ${n}`);return new c.LanguageModelToolCallPart(e,o,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,o)=>e.tools.length?{tools:e.tools.map(o),...e.toolRequired?{tool_choice:"required"}:{}}:{},G=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let o=0,n=0;for(let s of e.content){let t=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of t)o+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:w(r).length,n+=$(r)===void 0?0:1}return Math.ceil(o/4)+n*he},A=e=>e.map($).filter(o=>o!==void 0),E=e=>({results:e.filter(o=>o instanceof c.LanguageModelToolResultPart),calls:e.filter(o=>o instanceof c.LanguageModelToolCallPart),text:e.map(w).join(""),images:A(e)});function J(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant){let a=t.map(d=>({id:d.callId,type:"function",function:{name:d.name,arguments:JSON.stringify(d.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>A(a.content)),...l].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("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:o,stream:!0,...q(e,n=>({type:"function",function:{name:n.name,description:n.description,parameters:T(n)}}))}}var ye=e=>typeof e=="string"?e:(e??[]).filter(o=>o.type==="text").map(o=>o.text??"").join(""),ve=new Set(["stop","tool_calls","function_call"]);async function*F(e){let o=[],n;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 l=ye(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=o[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(n=r.finish_reason,!ve.has(n)))throw S(n)}}if(!n)throw P();for(let s of o.filter(Boolean))yield R(s.id,s.name,s.args)}var we=e=>{let o=e.map(w).join(""),n=A(e);return n.length?[...o?[{type:"input_text",text:o}]:[],...n.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:o};function H(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...t.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:we(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:o,stream:!0,store:!1,...q(e,n=>({type:"function",name:n.name,description:n.description,parameters:T(n),strict:!1}))}}async function*U(e){let o=[],n=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new c.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")o.push(R(t.item.call_id??"",t.item.name??"",t.item.arguments??""));else if(t.type==="response.completed")n=!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(!n)throw P();yield*o}var x=e=>e?[{type:"text",text:e}]:[],O=e=>e.flatMap(o=>{let n=D(o);return n?[{type:"image",source:{type:"base64",media_type:n.mimeType,data:n.data}}]:[]}),be={cache_control:{type:"ephemeral"}},N=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...be})};function B(e,o){let n=[],s=[];for(let t of e.messages){let{results:r,calls:l,text:i}=E(t.content);if(Number(t.role)===C){n.push(...x(i));continue}let a=t.role===c.LanguageModelChatMessageRole.Assistant?"assistant":"user",d=a==="assistant"?[...x(i),...l.map(g=>({type:"tool_use",id:g.callId,name:g.name,input:g.input}))]:[...r.map(g=>{let p=[...x(g.content.map(w).join("")),...O(g.content)];return{type:"tool_result",tool_use_id:g.callId,...p.length?{content:p}:{}}}),...x(i),...O(t.content)],u=s[s.length-1];if(d.length)u?.role===a?u.content.push(...d):s.push({role:a,content:d});else continue}return N(n),N(s[s.length-1]?.content??[]),{model:e.model,max_tokens:o,...n.length?{system:n}:{},messages:s,stream:!0,...e.tools.length?{tools:e.tools.map(t=>({name:t.name,description:t.description,input_schema:T(t)})),...e.toolRequired?{tool_choice:{type:"any"}}:{}}:{}}}var _e=new Set(["end_turn","tool_use","stop_sequence"]);async function*W(e){let o=new Map,n=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")o.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 c.LanguageModelTextPart(t.delta.text);else if(t.type==="content_block_delta"&&t.delta?.type==="input_json_delta"){let r=o.get(t.index??0);r&&(r.args+=t.delta.partial_json??"")}else{if(t.type==="message_delta"&&t.delta?.stop_reason&&!_e.has(t.delta.stop_reason))throw S(t.delta.stop_reason);if(t.type==="message_stop")n=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!n)throw P();for(let s of o.values())yield R(s.id,s.name,s.args)}async function*Y(e,o,n,s,t,r){let l=(p,f)=>fetch(`${e}${o}`,{method:"POST",headers:{...t,Authorization:`Bearer ${p}`,"Content-Type":"application/json"},body:JSON.stringify(f),signal:r}),i=await s.get(),a=await l(i,n);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i,n)),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 l(i,{...n,max_tokens:f}),!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`)}let d=new TextDecoder,u="";for await(let p of a.body){u+=d.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+d.decode());g&&g!=="[DONE]"&&(yield JSON.parse(g))}var K=e=>e.startsWith("data:")?e.slice(5).trim():"";var xe=["anthropic_messages","chat_completions","responses"],X="Databricks-Model-Provider-Service",ke=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(o=>{let n=e[o];return typeof n=="number"&&Number.isInteger(n)&&n>0?[[o,n]]:[]})),Q="ug-config.json",Z=e=>{let o={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(o=t)}catch{}let n=t=>o[t],s=n("models");return{workspace:String(n("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&xe.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(n("authCommand")??""),provider:typeof n("provider")=="string"?String(n("provider")):""}},ee=e=>`ucode/${M} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,k=class{constructor(o,n,s=()=>({})){this.settings=o;this.log=n;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:o,models:n,authCommand:s}=this.settings();return!o||!s||n.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):n.map(({id:t,name:r,family:l,api:i,maxInputTokens:a,maxOutputTokens:d})=>({id:t,api:i,name:r??t,family:l??t.split(".").pop()??t,version:M,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:d??32e3,capabilities:{toolCalling:!0,imageInput:i!=="chat_completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(o,n,s,t,r){let{workspace:l,authCommand:i,provider:a}=this.settings();this.tokens??=new _(i,m=>this.log(`Fetched a Databricks token (expires ${new Date(m).toISOString()})`));let d={model:o.id,messages:n,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(l,m,ne,f,{...p,...oe},u.signal),y=a?"/ai-gateway/openai/v1":"/ai-gateway/mlflow/v1",te={anthropic_messages:()=>W(v("/ai-gateway/anthropic/v1/messages",B(d,o.maxOutputTokens),{"anthropic-version":"2023-06-01"})),responses:()=>U(v(`${y}/responses`,H(d))),chat_completions:()=>F(v(`${y}/chat/completions`,J(d)))};try{for await(let m of te[o.api]())t.report(m)}catch(m){if(!r.isCancellationRequested)throw this.log(`${o.id} failed: ${m instanceof Error?m.message:String(m)}`),m}finally{g.dispose()}}async provideTokenCount(o,n){return G(n)}},Me=e=>{try{return(0,z.readFileSync)(e,"utf8")}catch{return}};function Ce(e){let o=h.window.createOutputChannel("Unity Gateway",{log:!0}),n=Z(Me((0,V.join)(e.extensionPath,Q))),s=()=>({"User-Agent":ee(h.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>n,r=>o.info(r),s);e.subscriptions.push(o,h.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,activate,readConfig,userAgent}); diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 3b0879694..a19c2dbbb 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -514,8 +514,13 @@ def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, @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"), + (["--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", diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 4f556e04d..1c9d8e8c8 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -134,18 +134,18 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider 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_messages: () => + anthropicMessagesReplies( + stream("/ai-gateway/anthropic/v1/messages", anthropicMessagesRequest(turn, model.maxOutputTokens), { + "anthropic-version": "2023-06-01", + }), + ), + responses: () => responsesReplies(stream(`${openai}/responses`, responsesRequest(turn))), + chat_completions: () => chatCompletionsReplies(stream(`${openai}/chat/completions`, chatCompletionsRequest(turn))), + }; try { - const replies: AsyncIterable = - model.api === "anthropic_messages" - ? anthropicMessagesReplies( - stream("/ai-gateway/anthropic/v1/messages", anthropicMessagesRequest(turn, model.maxOutputTokens), { - "anthropic-version": "2023-06-01", - }), - ) - : model.api === "responses" - ? responsesReplies(stream(`${openai}/responses`, responsesRequest(turn))) - : chatCompletionsReplies(stream(`${openai}/chat/completions`, chatCompletionsRequest(turn))); - 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/gateway.ts b/vscode-extension/src/gateway.ts index 4dba8f7a6..d07da86bd 100644 --- a/vscode-extension/src/gateway.ts +++ b/vscode-extension/src/gateway.ts @@ -9,26 +9,36 @@ export async function* streamEvents( headers: Record, signal: AbortSignal, ): AsyncGenerator { - const post = (token: string) => + 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 8f48baecc..900ced183 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -64,6 +64,26 @@ describe("streamEvents", () => { 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( From 2681edb207a363634b64405f9ff0004493ba2c66 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Fri, 30 Oct 2026 09:14:34 +0000 Subject: [PATCH 3/7] Name Copilot plainly in the provider-service messages Co-authored-by: Isaac --- src/ucode/agents/copilot.py | 2 +- tests/test_agent_copilot.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 8f96c55dd..ed979b368 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -380,7 +380,7 @@ def provider_service_catalog( 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 = "ug copilot" if tool == "copilot" else "VS Code" + agent = "Copilot" if tool == "copilot" else "VS Code" provider_type = service["provider_type"] named = list(service.get("targets") or []) listed: list[str] = [] 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): From e9c66117dd006f221ee1eb5dcee16824adb46873 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Fri, 30 Oct 2026 09:17:04 +0000 Subject: [PATCH 4/7] Write the catalog's wire name as the extension's api instead of a second set of names Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 13 +++---------- src/ucode/vscode/extension/out/extension.js | 4 ++-- tests/test_agent_vscode.py | 12 ++++++------ vscode-extension/README.md | 4 ++-- vscode-extension/src/extension.ts | 9 +++++---- vscode-extension/test/extension.test.ts | 20 ++++++++++---------- 6 files changed, 28 insertions(+), 34 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 0018590ce..befab1350 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -24,9 +24,6 @@ from xml.sax.saxutils import escape from ucode.agents.copilot import ( - WIRE_ANTHROPIC, - WIRE_COMPLETIONS, - WIRE_RESPONSES, CopilotModel, build_copilot_catalog, catalog_model_limits, @@ -49,12 +46,6 @@ # Each `code` command and the folder in the home directory it installs extensions under. CODE_COMMANDS = (("code", ".vscode"), ("code-insiders", ".vscode-insiders")) _CODE_TIMEOUT_S = 120 -# The extension's `api` for each wire API of Copilot's catalog. -_APIS = { - WIRE_ANTHROPIC: "anthropic_messages", - WIRE_RESPONSES: "responses", - WIRE_COMPLETIONS: "chat_completions", -} def _install_locations(command: str) -> list[Path]: @@ -181,7 +172,9 @@ def _model_entry(model: CopilotModel, name: str, *, external: bool) -> dict[str, "id": model.id, "name": name, "family": claude_model_name(model.id), - "api": _APIS[model.wire], + # The wire names of Copilot's catalog; the extension dispatches on them (see APIS in + # vscode-extension/src/extension.ts). + "api": model.wire, } # Models without known limits get the extension's 128K-in, 32K-out default. if limits := catalog_model_limits(model, external=external): diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index 92aa5d55f..d98a78f0b 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"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,o)=>{for(var n in o)b(e,n,{get:o[n],enumerable:!0})},j=(e,o,n,s)=>{if(o&&typeof o=="object"||typeof o=="function")for(let t of re(o))!ce.call(e,t)&&t!==n&&b(e,t,{get:()=>o[t],enumerable:!(s=ae(o,t))||s.enumerable});return e};var L=(e,o,n)=>(n=e!=null?se(ie(e)):{},j(o||!e||!e.__esModule?b(n,"default",{value:e,enumerable:!0}):n,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((o,n)=>{(0,I.exec)(e,{timeout:6e4,windowsHide:!0},(s,t,r)=>{s?n(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):o(String(t).trim())})}),fe=e=>{try{let{exp:o}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof o=="number")return o*1e3}catch{}return Date.now()+ue},_=class{constructor(o,n=()=>{}){this.command=o;this.onFetch=n}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(o){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==o?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(o){return this.pending??=me(o).then(n=>{if(!n)throw new Error(`Couldn't get a Databricks token: \`${o}\` printed nothing.`);return this.cached={value:n,expiresAt:fe(n)},this.onFetch(this.cached.expiresAt),n}).finally(()=>{this.pending=void 0}),this.pending}};var c=L(require("vscode")),C=3,he=1e3,w=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",D=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,$=e=>{let o=D(e);return o&&`data:${o.mimeType};base64,${o.data}`},T=e=>e.inputSchema??{type:"object",properties:{}},R=(e,o,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 ${o} with arguments that aren't a JSON object: ${n}`);return new c.LanguageModelToolCallPart(e,o,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,o)=>e.tools.length?{tools:e.tools.map(o),...e.toolRequired?{tool_choice:"required"}:{}}:{},G=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let o=0,n=0;for(let s of e.content){let t=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of t)o+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:w(r).length,n+=$(r)===void 0?0:1}return Math.ceil(o/4)+n*he},A=e=>e.map($).filter(o=>o!==void 0),E=e=>({results:e.filter(o=>o instanceof c.LanguageModelToolResultPart),calls:e.filter(o=>o instanceof c.LanguageModelToolCallPart),text:e.map(w).join(""),images:A(e)});function J(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant){let a=t.map(d=>({id:d.callId,type:"function",function:{name:d.name,arguments:JSON.stringify(d.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>A(a.content)),...l].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("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:o,stream:!0,...q(e,n=>({type:"function",function:{name:n.name,description:n.description,parameters:T(n)}}))}}var ye=e=>typeof e=="string"?e:(e??[]).filter(o=>o.type==="text").map(o=>o.text??"").join(""),ve=new Set(["stop","tool_calls","function_call"]);async function*F(e){let o=[],n;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 l=ye(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=o[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(n=r.finish_reason,!ve.has(n)))throw S(n)}}if(!n)throw P();for(let s of o.filter(Boolean))yield R(s.id,s.name,s.args)}var we=e=>{let o=e.map(w).join(""),n=A(e);return n.length?[...o?[{type:"input_text",text:o}]:[],...n.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:o};function H(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...t.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:we(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:o,stream:!0,store:!1,...q(e,n=>({type:"function",name:n.name,description:n.description,parameters:T(n),strict:!1}))}}async function*U(e){let o=[],n=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new c.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")o.push(R(t.item.call_id??"",t.item.name??"",t.item.arguments??""));else if(t.type==="response.completed")n=!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(!n)throw P();yield*o}var x=e=>e?[{type:"text",text:e}]:[],O=e=>e.flatMap(o=>{let n=D(o);return n?[{type:"image",source:{type:"base64",media_type:n.mimeType,data:n.data}}]:[]}),be={cache_control:{type:"ephemeral"}},N=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...be})};function B(e,o){let n=[],s=[];for(let t of e.messages){let{results:r,calls:l,text:i}=E(t.content);if(Number(t.role)===C){n.push(...x(i));continue}let a=t.role===c.LanguageModelChatMessageRole.Assistant?"assistant":"user",d=a==="assistant"?[...x(i),...l.map(g=>({type:"tool_use",id:g.callId,name:g.name,input:g.input}))]:[...r.map(g=>{let p=[...x(g.content.map(w).join("")),...O(g.content)];return{type:"tool_result",tool_use_id:g.callId,...p.length?{content:p}:{}}}),...x(i),...O(t.content)],u=s[s.length-1];if(d.length)u?.role===a?u.content.push(...d):s.push({role:a,content:d});else continue}return N(n),N(s[s.length-1]?.content??[]),{model:e.model,max_tokens:o,...n.length?{system:n}:{},messages:s,stream:!0,...e.tools.length?{tools:e.tools.map(t=>({name:t.name,description:t.description,input_schema:T(t)})),...e.toolRequired?{tool_choice:{type:"any"}}:{}}:{}}}var _e=new Set(["end_turn","tool_use","stop_sequence"]);async function*W(e){let o=new Map,n=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")o.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 c.LanguageModelTextPart(t.delta.text);else if(t.type==="content_block_delta"&&t.delta?.type==="input_json_delta"){let r=o.get(t.index??0);r&&(r.args+=t.delta.partial_json??"")}else{if(t.type==="message_delta"&&t.delta?.stop_reason&&!_e.has(t.delta.stop_reason))throw S(t.delta.stop_reason);if(t.type==="message_stop")n=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!n)throw P();for(let s of o.values())yield R(s.id,s.name,s.args)}async function*Y(e,o,n,s,t,r){let l=(p,f)=>fetch(`${e}${o}`,{method:"POST",headers:{...t,Authorization:`Bearer ${p}`,"Content-Type":"application/json"},body:JSON.stringify(f),signal:r}),i=await s.get(),a=await l(i,n);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i,n)),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 l(i,{...n,max_tokens:f}),!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`)}let d=new TextDecoder,u="";for await(let p of a.body){u+=d.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+d.decode());g&&g!=="[DONE]"&&(yield JSON.parse(g))}var K=e=>e.startsWith("data:")?e.slice(5).trim():"";var xe=["anthropic_messages","chat_completions","responses"],X="Databricks-Model-Provider-Service",ke=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(o=>{let n=e[o];return typeof n=="number"&&Number.isInteger(n)&&n>0?[[o,n]]:[]})),Q="ug-config.json",Z=e=>{let o={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(o=t)}catch{}let n=t=>o[t],s=n("models");return{workspace:String(n("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&xe.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(n("authCommand")??""),provider:typeof n("provider")=="string"?String(n("provider")):""}},ee=e=>`ucode/${M} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,k=class{constructor(o,n,s=()=>({})){this.settings=o;this.log=n;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:o,models:n,authCommand:s}=this.settings();return!o||!s||n.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):n.map(({id:t,name:r,family:l,api:i,maxInputTokens:a,maxOutputTokens:d})=>({id:t,api:i,name:r??t,family:l??t.split(".").pop()??t,version:M,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:d??32e3,capabilities:{toolCalling:!0,imageInput:i!=="chat_completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(o,n,s,t,r){let{workspace:l,authCommand:i,provider:a}=this.settings();this.tokens??=new _(i,m=>this.log(`Fetched a Databricks token (expires ${new Date(m).toISOString()})`));let d={model:o.id,messages:n,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(l,m,ne,f,{...p,...oe},u.signal),y=a?"/ai-gateway/openai/v1":"/ai-gateway/mlflow/v1",te={anthropic_messages:()=>W(v("/ai-gateway/anthropic/v1/messages",B(d,o.maxOutputTokens),{"anthropic-version":"2023-06-01"})),responses:()=>U(v(`${y}/responses`,H(d))),chat_completions:()=>F(v(`${y}/chat/completions`,J(d)))};try{for await(let m of te[o.api]())t.report(m)}catch(m){if(!r.isCancellationRequested)throw this.log(`${o.id} failed: ${m instanceof Error?m.message:String(m)}`),m}finally{g.dispose()}}async provideTokenCount(o,n){return G(n)}},Me=e=>{try{return(0,z.readFileSync)(e,"utf8")}catch{return}};function Ce(e){let o=h.window.createOutputChannel("Unity Gateway",{log:!0}),n=Z(Me((0,V.join)(e.extensionPath,Q))),s=()=>({"User-Agent":ee(h.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>n,r=>o.info(r),s);e.subscriptions.push(o,h.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,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,o)=>{for(var n in o)b(e,n,{get:o[n],enumerable:!0})},j=(e,o,n,s)=>{if(o&&typeof o=="object"||typeof o=="function")for(let t of re(o))!ce.call(e,t)&&t!==n&&b(e,t,{get:()=>o[t],enumerable:!(s=ae(o,t))||s.enumerable});return e};var L=(e,o,n)=>(n=e!=null?se(ie(e)):{},j(o||!e||!e.__esModule?b(n,"default",{value:e,enumerable:!0}):n,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((o,n)=>{(0,I.exec)(e,{timeout:6e4,windowsHide:!0},(s,t,r)=>{s?n(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):o(String(t).trim())})}),fe=e=>{try{let{exp:o}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof o=="number")return o*1e3}catch{}return Date.now()+ue},x=class{constructor(o,n=()=>{}){this.command=o;this.onFetch=n}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(o){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==o?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(o){return this.pending??=me(o).then(n=>{if(!n)throw new Error(`Couldn't get a Databricks token: \`${o}\` printed nothing.`);return this.cached={value:n,expiresAt:fe(n)},this.onFetch(this.cached.expiresAt),n}).finally(()=>{this.pending=void 0}),this.pending}};var c=L(require("vscode")),C=3,he=1e3,w=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",D=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,$=e=>{let o=D(e);return o&&`data:${o.mimeType};base64,${o.data}`},T=e=>e.inputSchema??{type:"object",properties:{}},R=(e,o,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 ${o} with arguments that aren't a JSON object: ${n}`);return new c.LanguageModelToolCallPart(e,o,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,o)=>e.tools.length?{tools:e.tools.map(o),...e.toolRequired?{tool_choice:"required"}:{}}:{},G=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let o=0,n=0;for(let s of e.content){let t=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of t)o+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:w(r).length,n+=$(r)===void 0?0:1}return Math.ceil(o/4)+n*he},A=e=>e.map($).filter(o=>o!==void 0),E=e=>({results:e.filter(o=>o instanceof c.LanguageModelToolResultPart),calls:e.filter(o=>o instanceof c.LanguageModelToolCallPart),text:e.map(w).join(""),images:A(e)});function J(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant){let a=t.map(d=>({id:d.callId,type:"function",function:{name:d.name,arguments:JSON.stringify(d.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>A(a.content)),...l].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("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:o,stream:!0,...q(e,n=>({type:"function",function:{name:n.name,description:n.description,parameters:T(n)}}))}}var ye=e=>typeof e=="string"?e:(e??[]).filter(o=>o.type==="text").map(o=>o.text??"").join(""),ve=new Set(["stop","tool_calls","function_call"]);async function*F(e){let o=[],n;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 l=ye(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=o[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(n=r.finish_reason,!ve.has(n)))throw S(n)}}if(!n)throw P();for(let s of o.filter(Boolean))yield R(s.id,s.name,s.args)}var we=e=>{let o=e.map(w).join(""),n=A(e);return n.length?[...o?[{type:"input_text",text:o}]:[],...n.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:o};function H(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...t.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:we(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:o,stream:!0,store:!1,...q(e,n=>({type:"function",name:n.name,description:n.description,parameters:T(n),strict:!1}))}}async function*U(e){let o=[],n=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new c.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")o.push(R(t.item.call_id??"",t.item.name??"",t.item.arguments??""));else if(t.type==="response.completed")n=!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(!n)throw P();yield*o}var _=e=>e?[{type:"text",text:e}]:[],O=e=>e.flatMap(o=>{let n=D(o);return n?[{type:"image",source:{type:"base64",media_type:n.mimeType,data:n.data}}]:[]}),be={cache_control:{type:"ephemeral"}},N=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...be})};function B(e,o){let n=[],s=[];for(let t of e.messages){let{results:r,calls:l,text:i}=E(t.content);if(Number(t.role)===C){n.push(..._(i));continue}let a=t.role===c.LanguageModelChatMessageRole.Assistant?"assistant":"user",d=a==="assistant"?[..._(i),...l.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}:{}}}),..._(i),...O(t.content)],u=s[s.length-1];if(d.length)u?.role===a?u.content.push(...d):s.push({role:a,content:d});else continue}return N(n),N(s[s.length-1]?.content??[]),{model:e.model,max_tokens:o,...n.length?{system:n}:{},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 o=new Map,n=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")o.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 c.LanguageModelTextPart(t.delta.text);else if(t.type==="content_block_delta"&&t.delta?.type==="input_json_delta"){let r=o.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")n=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!n)throw P();for(let s of o.values())yield R(s.id,s.name,s.args)}async function*Y(e,o,n,s,t,r){let l=(p,f)=>fetch(`${e}${o}`,{method:"POST",headers:{...t,Authorization:`Bearer ${p}`,"Content-Type":"application/json"},body:JSON.stringify(f),signal:r}),i=await s.get(),a=await l(i,n);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i,n)),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 l(i,{...n,max_tokens:f}),!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`)}let d=new TextDecoder,u="";for await(let p of a.body){u+=d.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+d.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(o=>{let n=e[o];return typeof n=="number"&&Number.isInteger(n)&&n>0?[[o,n]]:[]})),Q="ug-config.json",Z=e=>{let o={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(o=t)}catch{}let n=t=>o[t],s=n("models");return{workspace:String(n("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(n("authCommand")??""),provider:typeof n("provider")=="string"?String(n("provider")):""}},ee=e=>`ucode/${M} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,k=class{constructor(o,n,s=()=>({})){this.settings=o;this.log=n;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:o,models:n,authCommand:s}=this.settings();return!o||!s||n.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):n.map(({id:t,name:r,family:l,api:i,maxInputTokens:a,maxOutputTokens:d})=>({id:t,api:i,name:r??t,family:l??t.split(".").pop()??t,version:M,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:d??32e3,capabilities:{toolCalling:!0,imageInput:i!=="completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(o,n,s,t,r){let{workspace:l,authCommand:i,provider:a}=this.settings();this.tokens??=new x(i,m=>this.log(`Fetched a Databricks token (expires ${new Date(m).toISOString()})`));let d={model:o.id,messages:n,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(l,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(d,o.maxOutputTokens),{"anthropic-version":"2023-06-01"})),responses:()=>U(v(`${y}/responses`,H(d))),completions:()=>F(v(`${y}/chat/completions`,J(d)))};try{for await(let m of te[o.api]())t.report(m)}catch(m){if(!r.isCancellationRequested)throw this.log(`${o.id} failed: ${m instanceof Error?m.message:String(m)}`),m}finally{g.dispose()}}async provideTokenCount(o,n){return G(n)}},Me=e=>{try{return(0,z.readFileSync)(e,"utf8")}catch{return}};function Ce(e){let o=h.window.createOutputChannel("Unity Gateway",{log:!0}),n=Z(Me((0,V.join)(e.extensionPath,Q))),s=()=>({"User-Agent":ee(h.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>n,r=>o.info(r),s);e.subscriptions.push(o,h.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,activate,readConfig,userAgent}); diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index a19c2dbbb..248aa3474 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -110,7 +110,7 @@ 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", "family": "claude-sonnet-5-5", - "api": "anthropic_messages", + "api": "anthropic", "maxInputTokens": 872_000, "maxOutputTokens": 128_000, }, @@ -173,7 +173,7 @@ def test_a_provider_services_claude_targets_with_a_200k_cap_and_its_name( "id": "global.anthropic.claude-opus-4-8", "name": "Claude Opus 4.8", "family": "claude-opus-4-8", - "api": "anthropic_messages", + "api": "anthropic", "maxInputTokens": 168_000, "maxOutputTokens": 32_000, }, @@ -181,7 +181,7 @@ def test_a_provider_services_claude_targets_with_a_200k_cap_and_its_name( "id": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "name": "Claude Sonnet 4.5", "family": "claude-sonnet-4-5", - "api": "anthropic_messages", + "api": "anthropic", "maxInputTokens": 168_000, "maxOutputTokens": 32_000, }, @@ -212,7 +212,7 @@ def test_a_openai_provider_services_models_on_responses_and_chat_completions( 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": "chat_completions"}, + {"id": "gpt-4.1", "name": "GPT 4.1", "family": "gpt-4.1", "api": "completions"}, ] @pytest.mark.parametrize( @@ -265,7 +265,7 @@ def listing(workspace, token, *, parent_schema=None): "id": "main.models.claude-opus-5-5", "name": "Claude Opus 5.5", "family": "claude-opus-5-5", - "api": "anthropic_messages", + "api": "anthropic", "maxInputTokens": 168_000, "maxOutputTokens": 32_000, }, @@ -273,7 +273,7 @@ def listing(workspace, token, *, parent_schema=None): "id": "main.models.team-assistant", "name": "Team Assistant", "family": "team-assistant", - "api": "anthropic_messages", + "api": "anthropic", }, ] diff --git a/vscode-extension/README.md b/vscode-extension/README.md index 7ee043ca1..8abee96ef 100644 --- a/vscode-extension/README.md +++ b/vscode-extension/README.md @@ -5,8 +5,8 @@ mode. The extension calls the workspace's AI Gateway directly: Anthropic Message 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, family, `anthropic_messages`, `responses` or `chat_completions`, token -limits) and a command that prints a Databricks access token (`ug auth-token --host ...`), then installs +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 ` diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 1c9d8e8c8..4efd792e9 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -15,7 +15,8 @@ import { } from "./formats"; import { streamEvents } from "./gateway"; -const APIS = ["anthropic_messages", "chat_completions", "responses"] as const; +// 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; @@ -112,7 +113,7 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider // 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 !== "chat_completions" || id.includes("claude") }, + capabilities: { toolCalling: true, imageInput: api !== "completions" || id.includes("claude") }, })); } @@ -135,14 +136,14 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider // 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_messages: () => + anthropic: () => anthropicMessagesReplies( stream("/ai-gateway/anthropic/v1/messages", anthropicMessagesRequest(turn, model.maxOutputTokens), { "anthropic-version": "2023-06-01", }), ), responses: () => responsesReplies(stream(`${openai}/responses`, responsesRequest(turn))), - chat_completions: () => chatCompletionsReplies(stream(`${openai}/chat/completions`, chatCompletionsRequest(turn))), + completions: () => chatCompletionsReplies(stream(`${openai}/chat/completions`, chatCompletionsRequest(turn))), }; try { for await (const reply of apis[model.api]()) { diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 900ced183..af3d44225 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -116,12 +116,12 @@ describe("readConfig", () => { models: [ { id: "a", api: "responses" }, { id: "gpt-5.5", family: "gpt-5.5", api: "responses" }, - { id: "e", api: "anthropic_messages" }, + { 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)), { @@ -131,9 +131,9 @@ describe("readConfig", () => { models: [ { id: "a", api: "responses" }, { id: "gpt-5.5", family: "gpt-5.5", api: "responses" }, - { id: "e", api: "anthropic_messages" }, - { id: "c", name: "Model C", api: "chat_completions", maxInputTokens: 936_000 }, - { id: "d", api: "chat_completions" }, + { id: "e", api: "anthropic" }, + { id: "c", name: "Model C", api: "completions", maxInputTokens: 936_000 }, + { id: "d", api: "completions" }, ], }); }); @@ -152,9 +152,9 @@ describe("DatabricksModels", () => { authCommand: command, provider, models: [ - { id: "system.ai.claude-sonnet-5-5", name: "Claude Sonnet 5.5", api: "anthropic_messages", 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: "chat_completions" }, + { id: "gpt-4.1", family: "gpt-4.1", api: "completions" }, ], }); @@ -164,9 +164,9 @@ 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", "anthropic_messages", 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", "chat_completions", 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(), []); From 1f2519b6d604f7094d53110723b37db3095593bc Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Fri, 30 Oct 2026 10:25:41 +0000 Subject: [PATCH 5/7] Show a provider service's own model ids in the picker A service's targets are the provider's ids, which is what its users know it by, and its regional variants differ only by a prefix a readable label drops. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 8 +++++--- tests/test_agent_vscode.py | 10 +++++----- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index befab1350..72c354e00 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -207,13 +207,15 @@ def model_entries( ) -> 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.id if counts[label] > 1 else label, external=bool(provider or location) - ) + _model_entry(model, model.id if counts[label] > 1 else label, external=bool(location)) for model, label in labels.items() ] diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 248aa3474..27deea8cd 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -149,7 +149,7 @@ def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): "GPT 5.5", ] - def test_a_provider_services_claude_targets_with_a_200k_cap_and_its_name( + def test_a_provider_services_claude_targets_keep_their_ids_and_get_a_200k_cap( self, code, monkeypatch, capsys ): service = { @@ -171,7 +171,7 @@ def test_a_provider_services_claude_targets_with_a_200k_cap_and_its_name( assert code.config["models"] == [ { "id": "global.anthropic.claude-opus-4-8", - "name": "Claude Opus 4.8", + "name": "global.anthropic.claude-opus-4-8", "family": "claude-opus-4-8", "api": "anthropic", "maxInputTokens": 168_000, @@ -179,7 +179,7 @@ def test_a_provider_services_claude_targets_with_a_200k_cap_and_its_name( }, { "id": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", - "name": "Claude Sonnet 4.5", + "name": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "family": "claude-sonnet-4-5", "api": "anthropic", "maxInputTokens": 168_000, @@ -211,8 +211,8 @@ def test_a_openai_provider_services_models_on_responses_and_chat_completions( 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"}, + {"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( From 3a164b97869e1fb731a9ea8a4e1a153740add67e Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Fri, 30 Oct 2026 10:38:59 +0000 Subject: [PATCH 6/7] Send the Responses system prompt as instructions A system-role input item is refused outright by Microsoft Foundry, which made every turn through a Foundry service fail; OpenAI and Azure OpenAI accept the prompt either way. Co-authored-by: Isaac --- src/ucode/vscode/extension/out/extension.js | 6 ++++-- vscode-extension/src/formats.ts | 7 ++++++- vscode-extension/test/formats.test.ts | 2 +- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index d98a78f0b..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 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,o)=>{for(var n in o)b(e,n,{get:o[n],enumerable:!0})},j=(e,o,n,s)=>{if(o&&typeof o=="object"||typeof o=="function")for(let t of re(o))!ce.call(e,t)&&t!==n&&b(e,t,{get:()=>o[t],enumerable:!(s=ae(o,t))||s.enumerable});return e};var L=(e,o,n)=>(n=e!=null?se(ie(e)):{},j(o||!e||!e.__esModule?b(n,"default",{value:e,enumerable:!0}):n,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((o,n)=>{(0,I.exec)(e,{timeout:6e4,windowsHide:!0},(s,t,r)=>{s?n(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):o(String(t).trim())})}),fe=e=>{try{let{exp:o}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof o=="number")return o*1e3}catch{}return Date.now()+ue},x=class{constructor(o,n=()=>{}){this.command=o;this.onFetch=n}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(o){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==o?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(o){return this.pending??=me(o).then(n=>{if(!n)throw new Error(`Couldn't get a Databricks token: \`${o}\` printed nothing.`);return this.cached={value:n,expiresAt:fe(n)},this.onFetch(this.cached.expiresAt),n}).finally(()=>{this.pending=void 0}),this.pending}};var c=L(require("vscode")),C=3,he=1e3,w=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",D=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,$=e=>{let o=D(e);return o&&`data:${o.mimeType};base64,${o.data}`},T=e=>e.inputSchema??{type:"object",properties:{}},R=(e,o,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 ${o} with arguments that aren't a JSON object: ${n}`);return new c.LanguageModelToolCallPart(e,o,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,o)=>e.tools.length?{tools:e.tools.map(o),...e.toolRequired?{tool_choice:"required"}:{}}:{},G=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let o=0,n=0;for(let s of e.content){let t=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of t)o+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:w(r).length,n+=$(r)===void 0?0:1}return Math.ceil(o/4)+n*he},A=e=>e.map($).filter(o=>o!==void 0),E=e=>({results:e.filter(o=>o instanceof c.LanguageModelToolResultPart),calls:e.filter(o=>o instanceof c.LanguageModelToolCallPart),text:e.map(w).join(""),images:A(e)});function J(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant){let a=t.map(d=>({id:d.callId,type:"function",function:{name:d.name,arguments:JSON.stringify(d.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>A(a.content)),...l].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("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:o,stream:!0,...q(e,n=>({type:"function",function:{name:n.name,description:n.description,parameters:T(n)}}))}}var ye=e=>typeof e=="string"?e:(e??[]).filter(o=>o.type==="text").map(o=>o.text??"").join(""),ve=new Set(["stop","tool_calls","function_call"]);async function*F(e){let o=[],n;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 l=ye(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=o[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(n=r.finish_reason,!ve.has(n)))throw S(n)}}if(!n)throw P();for(let s of o.filter(Boolean))yield R(s.id,s.name,s.args)}var we=e=>{let o=e.map(w).join(""),n=A(e);return n.length?[...o?[{type:"input_text",text:o}]:[],...n.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:o};function H(e){let o=e.messages.flatMap(n=>{let{results:s,calls:t,text:r,images:l}=E(n.content);if(Number(n.role)===C)return[{role:"system",content:r}];if(n.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...t.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:we(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:o,stream:!0,store:!1,...q(e,n=>({type:"function",name:n.name,description:n.description,parameters:T(n),strict:!1}))}}async function*U(e){let o=[],n=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new c.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")o.push(R(t.item.call_id??"",t.item.name??"",t.item.arguments??""));else if(t.type==="response.completed")n=!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(!n)throw P();yield*o}var _=e=>e?[{type:"text",text:e}]:[],O=e=>e.flatMap(o=>{let n=D(o);return n?[{type:"image",source:{type:"base64",media_type:n.mimeType,data:n.data}}]:[]}),be={cache_control:{type:"ephemeral"}},N=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...be})};function B(e,o){let n=[],s=[];for(let t of e.messages){let{results:r,calls:l,text:i}=E(t.content);if(Number(t.role)===C){n.push(..._(i));continue}let a=t.role===c.LanguageModelChatMessageRole.Assistant?"assistant":"user",d=a==="assistant"?[..._(i),...l.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}:{}}}),..._(i),...O(t.content)],u=s[s.length-1];if(d.length)u?.role===a?u.content.push(...d):s.push({role:a,content:d});else continue}return N(n),N(s[s.length-1]?.content??[]),{model:e.model,max_tokens:o,...n.length?{system:n}:{},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 o=new Map,n=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")o.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 c.LanguageModelTextPart(t.delta.text);else if(t.type==="content_block_delta"&&t.delta?.type==="input_json_delta"){let r=o.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")n=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!n)throw P();for(let s of o.values())yield R(s.id,s.name,s.args)}async function*Y(e,o,n,s,t,r){let l=(p,f)=>fetch(`${e}${o}`,{method:"POST",headers:{...t,Authorization:`Bearer ${p}`,"Content-Type":"application/json"},body:JSON.stringify(f),signal:r}),i=await s.get(),a=await l(i,n);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i,n)),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 l(i,{...n,max_tokens:f}),!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`)}let d=new TextDecoder,u="";for await(let p of a.body){u+=d.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+d.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(o=>{let n=e[o];return typeof n=="number"&&Number.isInteger(n)&&n>0?[[o,n]]:[]})),Q="ug-config.json",Z=e=>{let o={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(o=t)}catch{}let n=t=>o[t],s=n("models");return{workspace:String(n("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(n("authCommand")??""),provider:typeof n("provider")=="string"?String(n("provider")):""}},ee=e=>`ucode/${M} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,k=class{constructor(o,n,s=()=>({})){this.settings=o;this.log=n;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:o,models:n,authCommand:s}=this.settings();return!o||!s||n.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):n.map(({id:t,name:r,family:l,api:i,maxInputTokens:a,maxOutputTokens:d})=>({id:t,api:i,name:r??t,family:l??t.split(".").pop()??t,version:M,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:d??32e3,capabilities:{toolCalling:!0,imageInput:i!=="completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(o,n,s,t,r){let{workspace:l,authCommand:i,provider:a}=this.settings();this.tokens??=new x(i,m=>this.log(`Fetched a Databricks token (expires ${new Date(m).toISOString()})`));let d={model:o.id,messages:n,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(l,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(d,o.maxOutputTokens),{"anthropic-version":"2023-06-01"})),responses:()=>U(v(`${y}/responses`,H(d))),completions:()=>F(v(`${y}/chat/completions`,J(d)))};try{for await(let m of te[o.api]())t.report(m)}catch(m){if(!r.isCancellationRequested)throw this.log(`${o.id} failed: ${m instanceof Error?m.message:String(m)}`),m}finally{g.dispose()}}async provideTokenCount(o,n){return G(n)}},Me=e=>{try{return(0,z.readFileSync)(e,"utf8")}catch{return}};function Ce(e){let o=h.window.createOutputChannel("Unity Gateway",{log:!0}),n=Z(Me((0,V.join)(e.extensionPath,Q))),s=()=>({"User-Agent":ee(h.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>n,r=>o.info(r),s);e.subscriptions.push(o,h.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,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/vscode-extension/src/formats.ts b/vscode-extension/src/formats.ts index b3fce3921..defbe7a26 100644 --- a/vscode-extension/src/formats.ts +++ b/vscode-extension/src/formats.ts @@ -194,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 [ @@ -216,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, diff --git a/vscode-extension/test/formats.test.ts b/vscode-extension/test/formats.test.ts index 067c05670..09f20604a 100644 --- a/vscode-extension/test/formats.test.ts +++ b/vscode-extension/test/formats.test.ts @@ -115,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"}' }, From da53dfb1e0867b3257923f2760c2c81397f93662 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Fri, 30 Oct 2026 10:47:05 +0000 Subject: [PATCH 7/7] Strip a slashed region prefix from a model name Azure names a regional deployment `eu/gpt-5.6-luna`, which no longer hid the model behind its region: such a model was classified for Chat Completions, which also switched off image input, and its family reached Copilot Chat unrecognisable. Co-authored-by: Isaac --- src/ucode/databricks.py | 3 ++- tests/test_databricks.py | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 4f00068c0..0fc2b9363 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3607,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/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):