diff --git a/.gitattributes b/.gitattributes index e21d1145a..47fd59c59 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ src/ucode/vscode/extension/out/extension.js linguist-generated=true +src/ucode/vscode/extension/package.json linguist-generated=true diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 5e7adfbe4..614f6c7a4 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -240,6 +240,11 @@ def _wire_from_name(name: str) -> str | None: return WIRE_COMPLETIONS if gpt else None +def uses_responses_api(model: str) -> bool: + """Whether Copilot calls ``model`` through the Responses API: GPT 5+ and o-series models.""" + return _wire_from_name(bare_model_name(model)) == WIRE_RESPONSES + + 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``. diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py new file mode 100644 index 000000000..bab6972e1 --- /dev/null +++ b/src/ucode/agents/vscode.py @@ -0,0 +1,272 @@ +"""VS Code: install ug's model extension so Copilot Chat can use the workspace's AI Gateway models. + +Configure-only, like Cursor (it stays out of `agents.__init__._MODULES`): `ug configure --agents +vscode` packs the extension that ships with ug into a vsix, together with a `ug-config.json` that +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. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import tempfile +import zipfile +from collections import Counter +from pathlib import Path +from typing import cast +from xml.sax.saxutils import escape + +from ucode.agents.copilot import ( + copilot_exclusion_reason, + copilot_model_limits, + uses_responses_api, +) +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.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 +from ucode.string_utils import model_picker_label +from ucode.ui import console, print_heading, print_kv, print_warning + +EXTENSION_DIR = Path(__file__).resolve().parent.parent / "vscode" / "extension" +EXTENSION_ID = "databricks.unity-gateway" +# Read by the extension from its install folder; see CONFIG_FILE in vscode-extension/src/extension.ts. +CONFIG_FILE = "ug-config.json" +# 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 + + +def _install_locations(command: str) -> list[Path]: + """Where VS Code's installer puts `command` when it isn't on PATH (by default on macOS).""" + insiders = command == "code-insiders" + system = current_os() + if system is OS.MACOS: + app = "Visual Studio Code - Insiders.app" if insiders else "Visual Studio Code.app" + return [ + root / app / "Contents/Resources/app/bin" / command + for root in (Path("/Applications"), Path.home() / "Applications") + ] + if system is OS.WINDOWS: + folder = "Microsoft VS Code Insiders" if insiders else "Microsoft VS Code" + roots = ( + [Path(os.environ["LOCALAPPDATA"]) / "Programs"] + if os.environ.get("LOCALAPPDATA") + else [] + ) + return [ + root / folder / "bin" / f"{command}.cmd" + for root in [*roots, windows_program_files_dir()] + ] + return [Path("/usr/bin") / command, Path("/snap/bin") / command] + + +def _find_command(command: str) -> Path | None: + if found := shutil.which(command): + return Path(found) + return next((path for path in _install_locations(command) if path.is_file()), None) + + +def find_code() -> Path | None: + """VS Code's `code` command, or Insiders' when only that is installed.""" + for command, _ in CODE_COMMANDS: + if found := _find_command(command): + return found + return None + + +def extensions_dir(folder: str) -> Path: + return Path.home() / folder / "extensions" + + +def _package() -> dict: + return json.loads((EXTENSION_DIR / "package.json").read_text(encoding="utf-8")) + + +def pack_vsix(directory: Path, config: dict[str, object]) -> Path: + """Zip the bundled extension and its ``config`` into a vsix in ``directory``, with no Node or + `vsce` needed.""" + package = _package() + manifest = f""" + + + + {escape(package["displayName"])} + {escape(package["description"])} + + + + + + + + + +""" + content_types = ( + '' + '' + '' + '' + '' + '' + "" + ) + vsix = directory / f"{package['name']}-{package['version']}.vsix" + with zipfile.ZipFile(vsix, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("[Content_Types].xml", content_types) + archive.writestr("extension.vsixmanifest", manifest) + for path in sorted(EXTENSION_DIR.rglob("*")): + if path.is_file(): + archive.write(path, f"extension/{path.relative_to(EXTENSION_DIR).as_posix()}") + archive.writestr(f"extension/{CONFIG_FILE}", json.dumps(config, indent=2) + "\n") + return vsix + + +def _run_code(code: Path, *args: str): + try: + return subprocess_cross_os.run( + [str(code), *args], capture_output=True, text=True, timeout=_CODE_TIMEOUT_S, check=False + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError(f"`{code} {args[0]}` failed: {exc}") from exc + + +def installed_version(code: Path) -> str | None: + result = _run_code(code, "--list-extensions", "--show-versions") + for line in (result.stdout or "").splitlines(): + name, _, version = line.strip().partition("@") + if name.lower() == EXTENSION_ID: + return version + return None + + +def install_extension(code: Path, config: dict[str, object]) -> str: + """Install the bundled extension with ``config`` over whatever is installed; returns what + happened.""" + # Always reinstall: the config changes with each configure, and a ug upgrade can change the + # bundle without changing its version. + version = _package()["version"] + with tempfile.TemporaryDirectory(prefix="ug-vscode-") as directory: + vsix = pack_vsix(Path(directory), config) + result = _run_code(code, "--install-extension", str(vsix), "--force") + if result.returncode != 0: + detail = (result.stderr or result.stdout or "").strip()[:300] + raise RuntimeError(f"`{code} --install-extension` failed: {detail}") + return f"installed {EXTENSION_ID} {version}" + + +def _model_entry(model: str, name: str) -> dict[str, object]: + entry: dict[str, object] = { + "id": model, + "name": name, + "api": "responses" if uses_responses_api(model) else "chat_completions", + } + # 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): + 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} + 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) + for model, label in labels.items() + ], reason + + +def auth_command(state: dict) -> str: + workspace = state["workspace"] + custom_oauth = state.get("custom_oauth") + if isinstance(custom_oauth, dict): + return build_custom_auth_shell_command(workspace, cast(CustomOAuthConfig, custom_oauth)) + return build_auth_shell_command( + workspace, state.get("profile"), use_pat=bool(state.get("use_pat")) + ) + + +def next_steps() -> list[str]: + """What to do in VS Code once the extension is installed, for this OS.""" + restart = ( + "quit it (Cmd+Q) and open it again" + if current_os() is OS.MACOS + else "close every VS Code window and open it again" + ) + return [ + # An open window keeps running the extension build it loaded, and configure reinstalls + # every time. A restart covers every open window at once. + f"1. If VS Code is open, {restart}.", + "2. Open Chat, click the model picker, and choose a Databricks model.", + ] + + +def configure(state: dict) -> bool: + """Install the extension with this workspace's setup. Returns False, after a warning, if any + step fails.""" + code = find_code() + if code is None: + print_warning( + "VS Code's `code` command wasn't found. Install VS Code (on macOS, also run " + "\"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: + # 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." + ) + return False + config = {"workspace": workspace, "authCommand": auth_command(state), "models": models} + try: + print_kv( + "VS Code extension", f"{install_extension(code, config)} with {len(models)} models" + ) + except RuntimeError as exc: + print_warning(f"{exc} Rerun `ug configure --agents vscode` to try again.") + return False + print_heading("Next, in VS Code:") + for line in next_steps(): + console.print(f" {line}", highlight=False) + return True + + +def revert() -> str: + """Uninstall the extension, and its config with it: "removed" or "unchanged".""" + removed = False + for command, folder in CODE_COMMANDS: + # Look on disk first, so revert runs `code` only for people who set VS Code up. + if not any(extensions_dir(folder).glob(f"{EXTENSION_ID}-*")): + continue + code = _find_command(command) + # VS Code can leave an uninstalled extension's folder behind until it next starts. + if code is None or installed_version(code) is None: + continue + result = _run_code(code, "--uninstall-extension", EXTENSION_ID) + if result.returncode == 0: + removed = True + else: + print_warning( + f"Couldn't uninstall {EXTENSION_ID}: {(result.stderr or '').strip()[:300]}" + ) + return "removed" if removed else "unchanged" diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 127dc0106..4746d22db 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1580,6 +1580,14 @@ def revert() -> int: claude_managed_result = claude_agent.revert_managed_settings() codex_managed_result = codex_agent.revert_managed_config() workspace_policy_result = revert_machine_policy() + from ucode.agents import vscode + + try: + vscode_result = vscode.revert() + except Exception as exc: # noqa: BLE001 -- Preserve other revert operations on failure + # VS Code's revert looks at the installed extension, not ug state, so a later run still finds it. + print_warning(f"Could not revert VS Code: {exc}. Fix it, then run ug revert again.") + vscode_result = "failed" results: dict[str, bool] = { # Antigravity's skills.json is shared with its app and IDE extensions, so revert @@ -1606,6 +1614,7 @@ def revert() -> int: print_kv("Codex OS-managed settings", codex_managed_result) print_kv("Unity Gateway workspace policy", workspace_policy_result) print_kv("Pi settings", "restored" if pi_settings_restored else "unchanged") + print_kv("VS Code extension", vscode_result) for client, spec in MCP_CLIENTS.items(): print_kv( f"{spec['display']} MCP config", @@ -4419,7 +4428,7 @@ def configure( str | None, typer.Option( "--agents", - help="Configure a comma-separated list of agents without prompting (e.g. claude,codex).", + help="Configure a comma-separated list of agents without prompting (e.g. claude,codex,vscode).", ), ] = None, workspace: Annotated[ @@ -4636,9 +4645,11 @@ def configure( # (which picks cursor up through MCP_ONLY_CLIENTS). If cursor is the # only agent, do a workspace-only configure so that a later `ug mcp # add` run has a current workspace to target. + # VS Code is configure-only too: it needs the workspace, then its own setup below. requested = [a.strip().lower() for a in agents.split(",") if a.strip()] wants_cursor = "cursor" in requested - model_agent_names = ",".join(a for a in requested if a != "cursor") + wants_vscode = "vscode" in requested + model_agent_names = ",".join(a for a in requested if a not in ("cursor", "vscode")) if model_agent_names: selected_tools = _parse_agents_option(model_agent_names) if workspace_entries is None: @@ -4652,8 +4663,8 @@ def configure( workspaces=workspace_entries, **skip_kwargs, ) - elif wants_cursor: - # Cursor-only: establish workspace state without the model picker. + elif wants_cursor or wants_vscode: + # Cursor/VS Code only: establish workspace state without the model picker. _configure_shared_workspace_states( workspace_entries or [_prompt_for_configuration(None)], tools=[], @@ -4664,6 +4675,17 @@ def configure( else: # Neither model agents nor cursor -> empty/invalid --agents list. _parse_agents_option(agents) + if wants_vscode and is_dry_run(): + # A dry run saves no workspace, so the state on disk isn't the one being configured. + print_note("Dry run: VS Code isn't changed.") + elif wants_vscode: + from ucode.agents import vscode + + # VS Code sends model traffic, so it gets the managed-workspace check other agents get. + state = load_state() + managed, _ = _read_policy_for_other_agent(state) + refuse_user_scope_over_managed_config(state.get("workspace"), managed) + vscode.configure(state) 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/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index 575c17279..c473cfd19 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var U=Object.create;var f=Object.defineProperty;var F=Object.getOwnPropertyDescriptor;var B=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,H=Object.prototype.hasOwnProperty;var K=(t,n)=>{for(var e in n)f(t,e,{get:n[e],enumerable:!0})},M=(t,n,e,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let i of B(n))!H.call(t,i)&&i!==e&&f(t,i,{get:()=>n[i],enumerable:!(s=F(n,i))||s.enumerable});return t};var C=(t,n,e)=>(e=t!=null?U(W(t)):{},M(n||!t||!t.__esModule?f(e,"default",{value:t,enumerable:!0}):e,t)),Y=t=>M(f({},"__esModule",{value:!0}),t);var se={};K(se,{DatabricksModels:()=>v,activate:()=>oe,readSettings:()=>$,userAgent:()=>q});module.exports=Y(se);var l=C(require("vscode"));var b="0.1.0";var T=require("node:child_process"),X=5*6e4,Q=60*6e4,V=t=>new Promise((n,e)=>{(0,T.exec)(t,{timeout:6e4,windowsHide:!0},(s,i,o)=>{s?e(new Error(`Couldn't get a Databricks token: \`${t}\` failed: ${String(o).trim().slice(0,500)||s.message}`)):n(String(i).trim())})}),Z=t=>{try{let{exp:n}=JSON.parse(Buffer.from(t.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+Q},h=class{constructor(n,e=()=>{}){this.command=n;this.onFetch=e}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-X>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??=V(n).then(e=>{if(!e)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:e,expiresAt:Z(e)},this.onFetch(this.cached.expiresAt),e}).finally(()=>{this.pending=void 0}),this.pending}};var c=C(require("vscode")),S=3,ee=1e3,y=t=>t instanceof c.LanguageModelTextPart?t.value:t instanceof c.LanguageModelDataPart&&!t.mimeType.startsWith("image/")?new TextDecoder().decode(t.data):t instanceof c.LanguageModelPromptTsxPart?JSON.stringify(t.value):"",_=t=>t instanceof c.LanguageModelDataPart&&t.mimeType.startsWith("image/")?`data:${t.mimeType};base64,${Buffer.from(t.data).toString("base64")}`:void 0,R=t=>t.inputSchema??{type:"object",properties:{}},L=(t,n,e)=>{let s;try{s=e.trim()?JSON.parse(e):{}}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: ${e}`);return new c.LanguageModelToolCallPart(t,n,s)},P=()=>new Error("The model stopped at its output token limit before finishing the response."),A=(t,n)=>t.tools.length?{tools:t.tools.map(n),...t.toolRequired?{tool_choice:"required"}:{}}:{},E=t=>{if(typeof t=="string")return Math.ceil(t.length/4);let n=0,e=0;for(let s of t.content){let i=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let o of i)n+=o instanceof c.LanguageModelToolCallPart?o.name.length+JSON.stringify(o.input).length:y(o).length,e+=_(o)===void 0?0:1}return Math.ceil(n/4)+e*ee},O=t=>({results:t.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:t.filter(n=>n instanceof c.LanguageModelToolCallPart),text:t.map(y).join(""),images:t.map(_).filter(n=>n!==void 0)});function I(t){let n=t.messages.flatMap(e=>{let{results:s,calls:i,text:o,images:r}=O(e.content);if(Number(e.role)===S)return[{role:"system",content:o}];if(e.role===c.LanguageModelChatMessageRole.Assistant){let a=i.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return o||a.length?[{role:"assistant",content:o||null,...a.length?{tool_calls:a}:{}}]:[]}let d=[...o?[{type:"text",text:o}]:[],...r.map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...d.length?[{role:"user",content:d}]:[]]});return{model:t.model,messages:n,stream:!0,...A(t,e=>({type:"function",function:{name:e.name,description:e.description,parameters:R(e)}}))}}var te=t=>typeof t=="string"?t:(t??[]).filter(n=>n.type==="text").map(n=>n.text??"").join("");async function*j(t){let n=[];for await(let e of t)for(let s of e.choices??[]){let i=te(s.delta?.content);i&&(yield new c.LanguageModelTextPart(i));for(let o of s.delta?.tool_calls??[]){let r=n[o.index??0]??={id:"",name:"",args:""};r.id=o.id||r.id,r.name=o.function?.name||r.name,r.args+=o.function?.arguments??""}if(s.finish_reason==="length")throw P()}for(let e of n.filter(Boolean))yield L(e.id,e.name,e.args)}function D(t){let n=t.messages.flatMap(e=>{let{results:s,calls:i,text:o,images:r}=O(e.content);if(Number(e.role)===S)return[{role:"system",content:o}];if(e.role===c.LanguageModelChatMessageRole.Assistant)return[...o?[{role:"assistant",content:o}]:[],...i.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let d=[...o?[{type:"input_text",text:o}]:[],...r.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:a.content.map(y).join("")})),...d.length?[{role:"user",content:d}]:[]]});return{model:t.model,input:n,stream:!0,store:!1,...A(t,e=>({type:"function",name:e.name,description:e.description,parameters:R(e),strict:!1}))}}async function*G(t){for await(let n of t){let e=n;if(e.type==="response.output_text.delta"&&e.delta)yield new c.LanguageModelTextPart(e.delta);else if(e.type==="response.output_item.done"&&e.item?.type==="function_call")yield L(e.item.call_id??"",e.item.name??"",e.item.arguments??"");else{if(e.type==="response.incomplete"&&e.response?.incomplete_details?.reason==="max_output_tokens")throw P();if(e.type==="response.failed"||e.type==="error")throw new Error(e.response?.error?.message??e.message??"The model response failed.")}}}async function*k(t,n,e,s,i,o){let r=m=>fetch(`${t}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...i,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(e),signal:o}),d=await s.get(),a=await r(d);if(a.status===401&&(await a.body?.cancel(),d=await s.renew(d),a=await r(d)),a.status===401)throw await a.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${t}\`, then retry.`);if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let x=g.split(` -`);g=x.pop()??"";for(let J of x){let w=N(J);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let u=N(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var N=t=>t.startsWith("data:")?t.slice(5).trim():"";var ne=t=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let e=t[n];return typeof e=="number"&&Number.isInteger(e)&&e>0?[[n,e]]:[]})),$=t=>{let n=t("models");return{workspace:String(t("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(e=>typeof e?.id=="string"&&e.id!==""&&(e.api==="chat_completions"||e.api==="responses")).map(e=>({id:e.id,api:e.api,...ne(e)})),authCommand:String(t("authCommand")??"")}},q=t=>`ucode/${b} githubcopilotchat${typeof t=="string"?`/${t}`:""}`,v=class{constructor(n,e,s=()=>({})){this.settings=n;this.log=e;this.headers=s}settings;log;headers;changed=new l.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:e,authCommand:s}=this.settings();return!n||!s||e.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):e.map(({id:i,api:o,maxInputTokens:r,maxOutputTokens:d})=>({id:i,api:o,name:i,family:i.split(".").pop()??i,version:b,detail:"Databricks",maxInputTokens:r??128e3,maxOutputTokens:d??32e3,capabilities:{toolCalling:!0,imageInput:o==="responses"||i.includes("claude")}}))}async provideLanguageModelChatResponse(n,e,s,i,o){let{workspace:r,authCommand:d}=this.settings();this.tokens??=new h(d,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:e,tools:s.tools??[],toolRequired:s.toolMode===l.LanguageModelChatToolMode.Required},p=new AbortController,g=o.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?G(k(r,"responses",D(a),this.tokens,this.headers(),p.signal)):j(k(r,"chat/completions",I(a),this.tokens,this.headers(),p.signal));for await(let m of u)i.report(m)}catch(u){if(!o.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,e){return E(e)}};function oe(t){let n=l.window.createOutputChannel("Unity Gateway",{log:!0}),e=()=>{let o=l.workspace.getConfiguration("unityGateway");return $(r=>o.get(r))},s=()=>({"User-Agent":q(l.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),i=new v(e,o=>n.info(o),s);t.subscriptions.push(n,l.lm.registerLanguageModelChatProvider("databricks-unity-gateway",i),l.commands.registerCommand("unityGateway.refreshModels",()=>i.refresh()),l.workspace.onDidChangeConfiguration(o=>{o.affectsConfiguration("unityGateway")&&i.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); +"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}); diff --git a/src/ucode/vscode/extension/package.json b/src/ucode/vscode/extension/package.json index 8bc404b46..aac96abc9 100644 --- a/src/ucode/vscode/extension/package.json +++ b/src/ucode/vscode/extension/package.json @@ -28,78 +28,6 @@ "displayName": "Databricks", "icon": "media/databricks.svg" } - ], - "configuration": { - "title": "Databricks Unity Gateway", - "properties": { - "unityGateway.workspace": { - "type": "string", - "default": "", - "description": "Databricks workspace URL, for example https://my-workspace.cloud.databricks.com. Set by `ug configure`.", - "scope": "application" - }, - "unityGateway.models": { - "type": "array", - "default": [], - "description": "Models to offer in the chat model picker. Set by `ug configure`.", - "items": { - "type": "object", - "required": [ - "id", - "api" - ], - "properties": { - "id": { - "type": "string", - "description": "Model service name, for example system.ai.claude-sonnet-5-5." - }, - "api": { - "type": "string", - "enum": [ - "chat_completions", - "responses" - ], - "description": "The AI Gateway API the model is called through." - }, - "maxInputTokens": { - "type": "integer", - "minimum": 1, - "description": "Prompt tokens the model accepts. Defaults to 128000." - }, - "maxOutputTokens": { - "type": "integer", - "minimum": 1, - "description": "Tokens the model can write in one reply. Defaults to 32000." - } - } - }, - "scope": "application" - }, - "unityGateway.authCommand": { - "type": "string", - "default": "", - "description": "Command that prints a Databricks access token. Set by `ug configure`.", - "scope": "application" - } - } - }, - "commands": [ - { - "command": "unityGateway.refreshModels", - "title": "Refresh Models", - "category": "Unity Gateway" - } ] - }, - "scripts": { - "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"const fs = require('fs'); fs.copyFileSync('package.json', '../src/ucode/vscode/extension/package.json'); fs.cpSync('media', '../src/ucode/vscode/extension/media', { recursive: true })\"", - "typecheck": "tsc", - "test": "esbuild test/*.test.ts --bundle --platform=node --format=cjs --target=node22 --outdir=.test-out --alias:vscode=./test/vscode.ts --log-level=warning && node --test .test-out/*.test.js" - }, - "devDependencies": { - "@types/node": "^22.20.3", - "@types/vscode": "1.125.0", - "esbuild": "^0.28.2", - "typescript": "^5.9.3" } } diff --git a/tests/conftest.py b/tests/conftest.py index 64e173255..a5450b505 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -41,6 +41,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch): from ucode.agents import antigravity as antigravity_mod from ucode.agents import claude as claude_mod from ucode.agents import codex as codex_mod + from ucode.agents import vscode as vscode_mod from ucode.agents.desktop import claude as desktop_mod state_dir = tmp_path / ".ucode" @@ -49,6 +50,12 @@ def _isolate_ucode_state(tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", state_dir) monkeypatch.setattr(codex_mod, "_otel_service_path", lambda: state_dir / "otel.service") monkeypatch.setattr(desktop_mod, "APP_DIR", state_dir) + # No test may run the developer's `code` or look at their installed VS Code extensions. + monkeypatch.setattr(vscode_mod, "find_code", lambda: None) + monkeypatch.setattr(vscode_mod, "_find_command", lambda command: None) + monkeypatch.setattr( + vscode_mod, "extensions_dir", lambda folder: tmp_path / folder / "extensions" + ) monkeypatch.setattr(desktop_auth_mod, "_AUTH_LOCK_PATH", state_dir / "claude-desktop-auth.lock") # Antigravity's customization root is shared with the developer's Antigravity app and CLI. agy_dir = tmp_path / ".gemini" / "config" diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py new file mode 100644 index 000000000..48c55402a --- /dev/null +++ b/tests/test_agent_vscode.py @@ -0,0 +1,367 @@ +"""Tests for `ug configure --agents vscode` and its revert.""" + +from __future__ import annotations + +import json +import subprocess +import zipfile +from pathlib import Path + +import pytest +from typer.testing import CliRunner + +from ucode import config_io +from ucode.agents import vscode +from ucode.cli import app +from ucode.managed_files import OS + +# The conftest fixture stubs these out for every test; the ones here exercise the real lookups. +REAL_FIND_CODE = vscode.find_code +REAL_FIND_COMMAND = vscode._find_command +REAL_RUN_CODE = vscode._run_code + +WS = "https://ws.cloud.databricks.com" +STATE = {"workspace": WS, "profile": "dev"} +VERSION = json.loads((vscode.EXTENSION_DIR / "package.json").read_text())["version"] + + +class FakeCode: + """Stands in for VS Code's `code` command: records calls and what an install unpacks.""" + + def __init__(self, extensions: Path) -> None: + self.extensions = extensions + self.fail: str | None = None + self.calls: list[tuple[str, ...]] = [] + self.paths: list[Path] = [] + self.vsix_entries: list[str] = [] + self.config: dict | None = None + + @property + def folder(self) -> Path: + return self.extensions / f"{vscode.EXTENSION_ID}-{VERSION}" + + def __call__(self, code: Path, *args: str) -> subprocess.CompletedProcess: + self.calls.append(args) + self.paths.append(code) + if args[0] == self.fail: + return subprocess.CompletedProcess(args, 1, "", "boom") + if args[0] == "--list-extensions": + listed = f"{vscode.EXTENSION_ID}@{VERSION}\n" if self.folder.exists() else "" + return subprocess.CompletedProcess(args, 0, "ms-python.python@2026.1.0\n" + listed, "") + if args[0] == "--install-extension": + with zipfile.ZipFile(args[1]) as vsix: + self.vsix_entries = vsix.namelist() + self.config = json.loads(vsix.read(f"extension/{vscode.CONFIG_FILE}")) + self.folder.mkdir(parents=True, exist_ok=True) + if args[0] == "--uninstall-extension": + self.folder.rmdir() + return subprocess.CompletedProcess(args, 0, "", "") + + +@pytest.fixture +def code(monkeypatch): + fake = FakeCode(vscode.extensions_dir(".vscode")) + monkeypatch.setattr(vscode, "find_code", lambda: Path("/fake/bin/code")) + 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( + vscode, + "discover_model_services", + lambda workspace, token: ( + {"sonnet": "system.ai.claude-sonnet-5-5"}, + ["system.ai.gpt-6-luna", "system.ai.gpt-oss-120b"], + ["system.ai.gemini-4-argon"], + ["system.ai.gpt-oss-120b", "system.ai.deepseek-v4-1-flash"], + None, + ), + ) + return fake + + +def _printed(capsys) -> str: + return " ".join(capsys.readouterr().out.split()) + + +class TestConfigure: + def test_installs_the_extension_with_the_workspace_and_its_claude_and_gpt_models( + self, code, capsys + ): + assert vscode.configure(STATE) + + assert code.calls[-1][0] == "--install-extension" + assert code.calls[-1][2] == "--force" + assert code.vsix_entries == [ + "[Content_Types].xml", + "extension.vsixmanifest", + "extension/media/databricks.svg", + "extension/out/extension.js", + "extension/package.json", + "extension/ug-config.json", + ] + assert code.config == { + "workspace": WS, + "authCommand": vscode.build_auth_shell_command(WS, "dev"), + "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.gpt-6-luna", "name": "GPT 6 Luna", "api": "responses"}, + ], + } + printed = _printed(capsys) + assert ( + f"VS Code extension: installed {vscode.EXTENSION_ID} {VERSION} with 2 models" in printed + ) + assert "Next, in VS Code: 1. If VS Code is open," in printed + + def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): + monkeypatch.setattr( + vscode, + "discover_model_services", + lambda workspace, token: ( + {}, + ["system.ai.gpt-6-luna", "main.models.gpt-6-luna", "system.ai.gpt-5-5"], + [], + [], + None, + ), + ) + + models, _ = vscode.model_entries(WS, "token") + + assert [model["name"] for model in models] == [ + "system.ai.gpt-6-luna", + "main.models.gpt-6-luna", + "GPT 5.5", + ] + + def test_reruns_reinstall_with_the_new_setup(self, code): + vscode.configure(STATE) + code.calls.clear() + + assert vscode.configure({**STATE, "profile": "other"}) + + assert "--profile other" in code.config["authCommand"] + assert [call[0] for call in code.calls] == ["--install-extension"] + + def test_warns_when_vscode_isnt_found(self, capsys): + assert not vscode.configure(STATE) + + assert "`code` command wasn't found" in _printed(capsys) + + @pytest.mark.parametrize("reason", ["the listing failed", None]) + def test_keeps_the_installed_setup_when_no_models_are_found( + self, code, monkeypatch, capsys, reason + ): + vscode.configure(STATE) + installed = code.config + code.calls.clear() + monkeypatch.setattr( + vscode, "discover_model_services", lambda workspace, token: ({}, [], [], [], reason) + ) + + assert not vscode.configure({**STATE, "workspace": "https://other.cloud.databricks.com"}) + + assert code.calls == [] + assert code.config == installed + assert f"No Claude or GPT models for VS Code ({reason or 'the workspace has none'})" in ( + _printed(capsys) + ) + + def test_reports_a_code_command_that_cant_run(self, code, monkeypatch, capsys): + def hang(*args, **kwargs): + raise subprocess.TimeoutExpired("code", 120) + + monkeypatch.setattr(vscode, "_run_code", REAL_RUN_CODE) + monkeypatch.setattr(vscode.subprocess_cross_os, "run", hang) + + assert not vscode.configure(STATE) + + assert f"`{Path('/fake/bin/code')} --install-extension` failed" in _printed(capsys) + + def test_reports_a_failed_install(self, code, capsys): + code.fail = "--install-extension" + + assert not vscode.configure(STATE) + + assert "--install-extension` failed: boom" in _printed(capsys) + + +class TestRevert: + def test_uninstalls_the_extension(self, code): + vscode.configure(STATE) + + assert vscode.revert() == "removed" + + assert code.calls[-1] == ("--uninstall-extension", vscode.EXTENSION_ID) + assert not code.folder.exists() + + def test_runs_no_code_command_when_the_extension_isnt_installed(self, code): + assert vscode.revert() == "unchanged" + + assert code.calls == [] + + def test_skips_a_folder_vscode_left_behind_after_an_uninstall(self, code): + (code.extensions / f"{vscode.EXTENSION_ID}-0.0.1").mkdir(parents=True) + + assert vscode.revert() == "unchanged" + + assert code.calls == [("--list-extensions", "--show-versions")] + + def test_uninstalls_through_the_insiders_command_for_an_insiders_install(self, code): + code.extensions = vscode.extensions_dir(".vscode-insiders") + vscode.configure(STATE) + + assert vscode.revert() == "removed" + + assert code.paths[-1] == Path("/fake/bin/code-insiders") + + +class TestFindCode: + @pytest.fixture(autouse=True) + def _real_lookup(self, monkeypatch): + monkeypatch.setattr(vscode, "find_code", REAL_FIND_CODE) + monkeypatch.setattr(vscode, "_find_command", REAL_FIND_COMMAND) + + def test_prefers_code_on_path(self, monkeypatch): + monkeypatch.setattr(vscode.shutil, "which", lambda name: f"/usr/local/bin/{name}") + + assert vscode.find_code() == Path("/usr/local/bin/code") + + def test_falls_back_to_insiders(self, monkeypatch): + monkeypatch.setattr( + vscode.shutil, + "which", + lambda name: "/opt/bin/code-insiders" if name == "code-insiders" else None, + ) + + assert vscode.find_code() == Path("/opt/bin/code-insiders") + + def test_looks_in_both_macos_application_folders(self, monkeypatch, tmp_path): + monkeypatch.setattr(vscode, "current_os", lambda: OS.MACOS) + monkeypatch.setattr(vscode.Path, "home", lambda: tmp_path) + app = "Visual Studio Code.app/Contents/Resources/app/bin/code" + + assert vscode._install_locations("code") == [ + Path("/Applications") / app, + tmp_path / "Applications" / app, + ] + + def test_falls_back_to_the_first_install_location_that_exists(self, monkeypatch, tmp_path): + monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + installed = tmp_path / "bin/code" + installed.parent.mkdir() + installed.touch() + monkeypatch.setattr( + vscode, + "_install_locations", + lambda command: [tmp_path / "missing/code", installed] if command == "code" else [], + ) + + assert vscode.find_code() == installed + + def test_finds_the_windows_user_install(self, monkeypatch, tmp_path): + monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + monkeypatch.setattr(vscode, "current_os", lambda: OS.WINDOWS) + monkeypatch.setenv("LOCALAPPDATA", str(tmp_path)) + code_cmd = tmp_path / "Programs/Microsoft VS Code/bin/code.cmd" + code_cmd.parent.mkdir(parents=True) + code_cmd.touch() + + assert vscode.find_code() == code_cmd + + def test_none_when_vscode_isnt_installed(self, monkeypatch): + monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + monkeypatch.setattr(vscode, "_install_locations", lambda command: []) + + assert vscode.find_code() is None + + +class TestNextSteps: + @pytest.mark.parametrize( + ("system", "restart"), + [ + (OS.MACOS, "If VS Code is open, quit it (Cmd+Q) and open it again."), + (OS.WINDOWS, "If VS Code is open, close every VS Code window and open it again."), + (OS.LINUX, "If VS Code is open, close every VS Code window and open it again."), + ], + ) + def test_says_how_to_restart_on_the_os(self, monkeypatch, system, restart): + monkeypatch.setattr(vscode, "current_os", lambda: system) + + assert vscode.next_steps() == [ + f"1. {restart}", + "2. Open Chat, click the model picker, and choose a Databricks model.", + ] + + +class TestAuthCommand: + def test_pat_and_custom_oauth(self): + assert "--use-pat" in vscode.auth_command({**STATE, "use_pat": True}) + custom = { + "client_id": "c", + "redirect_url": "http://localhost:8020", + "scopes": "all-apis", + "profile": "custom", + } + assert "--profile custom" in vscode.auth_command({"workspace": WS, "custom_oauth": custom}) + + +def test_configure_agents_vscode_refuses_a_workspace_the_machine_isnt_managed_for(monkeypatch): + selected: list[str] = [] + monkeypatch.setattr( + "ucode.cli._configure_shared_workspace_states", lambda *args, **kwargs: [STATE] + ) + monkeypatch.setattr("ucode.cli.load_state", lambda: STATE) + monkeypatch.setattr( + "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")) + + result = CliRunner().invoke(app, ["configure", "--agents", "vscode", "--workspace", WS]) + + assert result.exit_code == 1 + assert "This machine is managed for https://managed.cloud.databricks.com" in " ".join( + result.output.split() + ) + assert selected == ["https://managed.cloud.databricks.com"] + + +@pytest.mark.parametrize("dry_run", [False, True]) +def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, dry_run): + # `--dry-run` sets a module-wide flag; this puts it back so later tests write files again. + monkeypatch.setattr(config_io, "_dry_run", False) + configured: list[dict] = [] + shared: list[dict] = [] + monkeypatch.setattr( + "ucode.cli._configure_shared_workspace_states", + lambda *args, **kwargs: shared.append(kwargs) or [], + ) + monkeypatch.setattr( + "ucode.cli.configure_workspace_command", + 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) + + args = [ + "configure", + "--agents", + "vscode", + "--workspace", + WS, + *(["--dry-run"] if dry_run else []), + ] + result = CliRunner().invoke(app, args) + + 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]) + assert ("VS Code isn't changed" in result.output) == dry_run diff --git a/tests/test_cli.py b/tests/test_cli.py index c15012a7e..f460fc67b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3971,6 +3971,28 @@ def test_desktop_revert_failure_is_actionable(self): assert "ug state retained" in result.output assert "Fix the Claude Desktop error and run ug revert again" in result.output + def test_vscode_revert_failure_is_reported_and_the_rest_still_reverts(self): + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.agents.desktop.setup.revert_desktop", return_value="unchanged"), + patch("ucode.cli.revert_mcp_configs", return_value={}), + patch("ucode.cli.claude_agent.revert_managed_settings", return_value="unchanged"), + patch("ucode.cli.codex_agent.revert_managed_config", return_value="unchanged"), + patch("ucode.cli.revert_machine_policy", return_value="unchanged"), + patch("ucode.cli.restore_file", return_value=False) as restore_file, + patch("ucode.cli.revert_legacy_shared_config", return_value=False), + patch("ucode.agents.vscode.revert", side_effect=RuntimeError("code timed out")), + patch("ucode.cli.clear_state") as clear_state, + ): + result = runner.invoke(app, ["revert"]) + + output = " ".join(result.output.split()) + assert result.exit_code == 0, result.output + assert "Could not revert VS Code: code timed out" in output + assert "VS Code extension: failed" in output + assert restore_file.called + clear_state.assert_called_once() + @pytest.mark.parametrize("args", [[], ["claude"], ["codex"]]) def test_unknown_routing_config_does_not_block_launch(monkeypatch, args): diff --git a/vscode-extension/README.md b/vscode-extension/README.md index d49c502a9..b79ed3ea4 100644 --- a/vscode-extension/README.md +++ b/vscode-extension/README.md @@ -1,20 +1,16 @@ # 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 and open -models, Responses for GPT-5 and later. +mode. The extension calls the workspace's AI Gateway directly: Chat Completions for Claude, Responses +for GPT-5 and later. -`ug configure` installs it and writes its settings: - -| Setting | Value | -|---|---| -| `unityGateway.workspace` | The workspace URL | -| `unityGateway.models` | `[{"id": "", "api": "chat_completions" \| "responses"}]` | -| `unityGateway.authCommand` | A command that prints a Databricks access token (`ug auth-token --host ...`) | +`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 token is kept in memory until it nears expiry, and fetched again when the gateway answers 401. -Run **Unity Gateway: Refresh Models** after changing the settings by hand. Token fetches and failed -requests are logged to the **Unity Gateway** output channel. +Token fetches and failed requests are logged to the **Unity Gateway** output channel. ## Development diff --git a/vscode-extension/media/databricks.svg b/vscode-extension/media/databricks.svg deleted file mode 100644 index b7bd91e12..000000000 --- a/vscode-extension/media/databricks.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/vscode-extension/package.json b/vscode-extension/package.json index 8bc404b46..dc5d88801 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -28,71 +28,10 @@ "displayName": "Databricks", "icon": "media/databricks.svg" } - ], - "configuration": { - "title": "Databricks Unity Gateway", - "properties": { - "unityGateway.workspace": { - "type": "string", - "default": "", - "description": "Databricks workspace URL, for example https://my-workspace.cloud.databricks.com. Set by `ug configure`.", - "scope": "application" - }, - "unityGateway.models": { - "type": "array", - "default": [], - "description": "Models to offer in the chat model picker. Set by `ug configure`.", - "items": { - "type": "object", - "required": [ - "id", - "api" - ], - "properties": { - "id": { - "type": "string", - "description": "Model service name, for example system.ai.claude-sonnet-5-5." - }, - "api": { - "type": "string", - "enum": [ - "chat_completions", - "responses" - ], - "description": "The AI Gateway API the model is called through." - }, - "maxInputTokens": { - "type": "integer", - "minimum": 1, - "description": "Prompt tokens the model accepts. Defaults to 128000." - }, - "maxOutputTokens": { - "type": "integer", - "minimum": 1, - "description": "Tokens the model can write in one reply. Defaults to 32000." - } - } - }, - "scope": "application" - }, - "unityGateway.authCommand": { - "type": "string", - "default": "", - "description": "Command that prints a Databricks access token. Set by `ug configure`.", - "scope": "application" - } - } - }, - "commands": [ - { - "command": "unityGateway.refreshModels", - "title": "Refresh Models", - "category": "Unity Gateway" - } ] }, "scripts": { - "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"const fs = require('fs'); fs.copyFileSync('package.json', '../src/ucode/vscode/extension/package.json'); fs.cpSync('media', '../src/ucode/vscode/extension/media', { recursive: true })\"", + "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"const m = require('./package.json'); delete m.scripts; delete m.devDependencies; require('fs').writeFileSync('../src/ucode/vscode/extension/package.json', JSON.stringify(m, null, 2) + '\\n')\"", "typecheck": "tsc", "test": "esbuild test/*.test.ts --bundle --platform=node --format=cjs --target=node22 --outdir=.test-out --alias:vscode=./test/vscode.ts --log-level=warning && node --test .test-out/*.test.js" }, diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 991e27485..2ba3da005 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -1,3 +1,5 @@ +import { readFileSync } from "node:fs"; +import { join } from "node:path"; import * as vscode from "vscode"; import { version } from "../package.json"; import { Tokens } from "./auth"; @@ -6,6 +8,7 @@ import { streamEvents } from "./gateway"; export interface ModelSetting { readonly id: string; + readonly name?: string; readonly api: "chat_completions" | "responses"; readonly maxInputTokens?: number; readonly maxOutputTokens?: number; @@ -25,7 +28,21 @@ const tokenLimits = (model: ModelSetting): Partial => }), ); -export const readSettings = (get: (key: string) => unknown): Settings => { +/** `ug configure` writes this file into the extension it packs, so each install carries its own setup. */ +export const CONFIG_FILE = "ug-config.json"; + +/** Reads the config file's text; a missing or malformed file gives no models. */ +export const readConfig = (text: string | undefined): Settings => { + let config: Record = {}; + try { + const parsed: unknown = JSON.parse(text ?? "{}"); + if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) { + config = parsed as Record; + } + } catch { + // Treated like a missing file. + } + const get = (key: string) => config[key]; const models = get("models"); return { workspace: String(get("workspace") ?? "").replace(/\/+$/, ""), @@ -34,7 +51,12 @@ export const readSettings = (get: (key: string) => unknown): Settings => { (model): model is ModelSetting => typeof model?.id === "string" && model.id !== "" && (model.api === "chat_completions" || model.api === "responses"), ) - .map((model) => ({ id: model.id, api: model.api, ...tokenLimits(model) })), + .map((model) => ({ + id: model.id, + ...(typeof model.name === "string" && model.name.trim() ? { name: model.name.trim() } : {}), + api: model.api, + ...tokenLimits(model), + })), authCommand: String(get("authCommand") ?? ""), }; }; @@ -46,8 +68,6 @@ export const userAgent = (copilotChatVersion: unknown): string => type Model = vscode.LanguageModelChatInformation & { readonly api: ModelSetting["api"] }; export class DatabricksModels implements vscode.LanguageModelChatProvider { - private readonly changed = new vscode.EventEmitter(); - readonly onDidChangeLanguageModelChatInformation = this.changed.event; private tokens: Tokens | undefined; constructor( @@ -56,22 +76,16 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider private readonly headers: () => Record = () => ({}), ) {} - /** Re-reads the settings: drops the cached token and asks VS Code to list the models again. */ - refresh(): void { - this.tokens = undefined; - this.changed.fire(); - } - async provideLanguageModelChatInformation(): Promise { const { workspace, models, authCommand } = this.settings(); if (!workspace || !authCommand || models.length === 0) { this.log("No Databricks models: run `ug configure --agents vscode` to set them up."); return []; } - return models.map(({ id, api, maxInputTokens, maxOutputTokens }) => ({ + return models.map(({ id, name, api, maxInputTokens, maxOutputTokens }) => ({ id, api, - name: id, + 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, version, @@ -118,22 +132,19 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider } } +const readConfigFile = (path: string): string | undefined => { + try { + return readFileSync(path, "utf8"); + } catch { + return undefined; + } +}; + export function activate(context: vscode.ExtensionContext): void { const output = vscode.window.createOutputChannel("Unity Gateway", { log: true }); - const settings = () => { - const configuration = vscode.workspace.getConfiguration("unityGateway"); - return readSettings((key) => configuration.get(key)); - }; + // Read once: ug changes the file only by reinstalling the extension, which takes a restart to load. + const settings = readConfig(readConfigFile(join(context.extensionPath, CONFIG_FILE))); const headers = () => ({ "User-Agent": userAgent(vscode.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version) }); - const provider = new DatabricksModels(settings, (line) => output.info(line), headers); - context.subscriptions.push( - output, - vscode.lm.registerLanguageModelChatProvider("databricks-unity-gateway", provider), - vscode.commands.registerCommand("unityGateway.refreshModels", () => provider.refresh()), - vscode.workspace.onDidChangeConfiguration((event) => { - if (event.affectsConfiguration("unityGateway")) { - provider.refresh(); - } - }), - ); + const provider = new DatabricksModels(() => settings, (line) => output.info(line), headers); + context.subscriptions.push(output, vscode.lm.registerLanguageModelChatProvider("databricks-unity-gateway", provider)); } diff --git a/vscode-extension/src/formats.ts b/vscode-extension/src/formats.ts index d2c61f7e9..3e2be406f 100644 --- a/vscode-extension/src/formats.ts +++ b/vscode-extension/src/formats.ts @@ -48,7 +48,16 @@ const toolCall = (callId: string, name: string, args: string): vscode.LanguageMo return new vscode.LanguageModelToolCallPart(callId, name, input); }; -const outputLimit = () => new Error("The model stopped at its output token limit before finishing the response."); +/** Why the model stopped short, from a Responses `incomplete_details.reason` or a Chat Completions `finish_reason`. */ +const stoppedEarly = (reason: string | undefined): Error => + reason === "max_output_tokens" || reason === "length" + ? new Error("The model stopped at its output token limit before finishing the response.") + : reason === "content_filter" + ? 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"}).`); + +// A dropped connection ends the stream like a finished answer does, so only a final event counts as done. +const cutOff = () => new Error("The response ended before the model finished. Try again."); const toolOptions = (turn: Turn, toTool: (tool: vscode.LanguageModelChatTool) => object): object => turn.tools.length ? { tools: turn.tools.map(toTool), ...(turn.toolRequired ? { tool_choice: "required" } : {}) } : {}; @@ -70,12 +79,14 @@ export const estimateTokens = (input: string | vscode.LanguageModelChatRequestMe return Math.ceil(characters / 4) + images * IMAGE_TOKENS; }; +const imagesOf = (parts: readonly unknown[]): string[] => parts.map(imageUrl).filter((url) => url !== undefined); + /** Splits a message into its tool results, tool calls, other text, and images. */ const sortParts = (content: readonly unknown[]) => ({ results: content.filter((part) => part instanceof vscode.LanguageModelToolResultPart), calls: content.filter((part) => part instanceof vscode.LanguageModelToolCallPart), text: content.map(textOf).join(""), - images: content.map(imageUrl).filter((url) => url !== undefined), + images: imagesOf(content), }); export function chatCompletionsRequest(turn: Turn): object { @@ -92,9 +103,10 @@ export function chatCompletionsRequest(turn: Turn): object { })); return text || toolCalls.length ? [{ role: "assistant", content: text || null, ...(toolCalls.length ? { tool_calls: toolCalls } : {}) }] : []; } + // The gateway rejects images in a tool message, so a tool's images follow its result in the user message. const content = [ ...(text ? [{ type: "text", text }] : []), - ...images.map((url) => ({ type: "image_url", image_url: { url } })), + ...[...results.flatMap((result) => imagesOf(result.content)), ...images].map((url) => ({ type: "image_url", image_url: { url } })), ]; return [ ...results.map((result) => ({ role: "tool", tool_call_id: result.callId, content: result.content.map(textOf).join("") })), @@ -112,6 +124,7 @@ export function chatCompletionsRequest(turn: Turn): object { type ChatContent = string | { type?: string; text?: string }[] | null; interface ChatChunk { + error?: string | { message?: string }; choices?: { delta?: { content?: ChatContent; tool_calls?: { index?: number; id?: string; function?: { name?: string; arguments?: string } }[] }; finish_reason?: string | null; @@ -128,9 +141,16 @@ const deltaText = (content: ChatContent | undefined): string => .map((block) => block.text ?? "") .join(""); +const CHAT_FINISHED = new Set(["stop", "tool_calls", "function_call"]); + export async function* chatCompletionsReplies(events: AsyncIterable): AsyncGenerator { const calls: { id: string; name: string; args: string }[] = []; + let finish: string | undefined; for await (const event of events) { + const { error } = event as ChatChunk; + if (error) { + throw new Error(typeof error === "string" ? error : (error.message ?? JSON.stringify(error))); + } for (const choice of (event as ChatChunk).choices ?? []) { const text = deltaText(choice.delta?.content); if (text) { @@ -142,16 +162,30 @@ export async function* chatCompletionsReplies(events: AsyncIterable): A call.name = fragment.function?.name || call.name; call.args += fragment.function?.arguments ?? ""; } - if (choice.finish_reason === "length") { - throw outputLimit(); + if (choice.finish_reason) { + finish = choice.finish_reason; + if (!CHAT_FINISHED.has(finish)) { + throw stoppedEarly(finish); + } } } } + if (!finish) { + throw cutOff(); + } for (const call of calls.filter(Boolean)) { yield toolCall(call.id, call.name, call.args); } } +const functionOutput = (content: readonly unknown[]): string | object[] => { + const text = content.map(textOf).join(""); + const images = imagesOf(content); + return images.length + ? [...(text ? [{ type: "input_text", text }] : []), ...images.map((url) => ({ type: "input_image", image_url: url, detail: "auto" }))] + : text; +}; + // Stateless: the gateway keeps no response state, so every turn resends the whole conversation. export function responsesRequest(turn: Turn): object { const input = turn.messages.flatMap((message): object[] => { @@ -170,7 +204,7 @@ export function responsesRequest(turn: Turn): object { ...images.map((url) => ({ type: "input_image", image_url: url, detail: "auto" })), ]; return [ - ...results.map((result) => ({ type: "function_call_output", call_id: result.callId, output: result.content.map(textOf).join("") })), + ...results.map((result) => ({ type: "function_call_output", call_id: result.callId, output: functionOutput(result.content) })), ...(content.length ? [{ role: "user", content }] : []), ]; }); @@ -192,16 +226,25 @@ interface ResponseEvent { } export async function* responsesReplies(events: AsyncIterable): AsyncGenerator { + // Held back until the response completes, so a cut-off stream runs no tool. + const calls: vscode.LanguageModelToolCallPart[] = []; + let completed = false; for await (const raw of events) { const event = raw as ResponseEvent; if (event.type === "response.output_text.delta" && event.delta) { yield new vscode.LanguageModelTextPart(event.delta); } else if (event.type === "response.output_item.done" && event.item?.type === "function_call") { - yield toolCall(event.item.call_id ?? "", event.item.name ?? "", event.item.arguments ?? ""); - } else if (event.type === "response.incomplete" && event.response?.incomplete_details?.reason === "max_output_tokens") { - throw outputLimit(); + calls.push(toolCall(event.item.call_id ?? "", event.item.name ?? "", event.item.arguments ?? "")); + } else if (event.type === "response.completed") { + completed = true; + } else if (event.type === "response.incomplete") { + throw stoppedEarly(event.response?.incomplete_details?.reason); } else if (event.type === "response.failed" || event.type === "error") { throw new Error(event.response?.error?.message ?? event.message ?? "The model response failed."); } } + if (!completed) { + throw cutOff(); + } + yield* calls; } diff --git a/vscode-extension/src/gateway.ts b/vscode-extension/src/gateway.ts index 96d5eafdd..5db9b7959 100644 --- a/vscode-extension/src/gateway.ts +++ b/vscode-extension/src/gateway.ts @@ -25,7 +25,7 @@ export async function* streamEvents( } if (response.status === 401) { await response.body?.cancel(); - throw new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${workspace}\`, then retry.`); + 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)}`); diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 5e94cfcf4..537380ec8 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, readSettings, userAgent, type Settings } from "../src/extension"; +import { DatabricksModels, 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,10 @@ 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; body: Record }[] = []; + const requests: { url: string; authorization: string; userAgent: string; body: Record }[] = []; globalThis.fetch = (async (url: string, init: RequestInit) => { const headers = new Headers(init.headers); - requests.push({ url, authorization: headers.get("authorization") ?? "", body: JSON.parse(String(init.body)) }); + requests.push({ url, authorization: headers.get("authorization") ?? "", userAgent: headers.get("user-agent") ?? "", 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); @@ -53,7 +53,7 @@ describe("streamEvents", () => { requests = stubFetch([401, ""]); await rejectsWith( collect(streamEvents("https://ws", "responses", {}, new Tokens(tokenCommand(["a", "b"]).command), {}, new AbortController().signal)), - "Run `databricks auth login --host https://ws`", + "Run `ug configure --agents vscode` to sign in again", ); assert.equal(requests.length, 2); }); @@ -67,17 +67,10 @@ describe("streamEvents", () => { }); }); -describe("settings", () => { - // Workspace settings mustn't redirect the token or choose the command that prints it. - it("are read from user settings only", () => { - const properties = Object.values(packageJson.contributes.configuration.properties); - assert.deepEqual(properties.map((property) => property.scope), ["application", "application", "application"]); - }); -}); - describe("provider", () => { + // The icon lives only in the folder ug ships, next to the built bundle. it("ships the icon it names", () => { - assert.ok(packageJson.contributes.languageModelChatProviders.every((provider) => existsSync(provider.icon))); + assert.ok(packageJson.contributes.languageModelChatProviders.every((provider) => existsSync(`../src/ucode/vscode/extension/${provider.icon}`))); }); }); @@ -88,9 +81,9 @@ describe("userAgent", () => { }); }); -describe("readSettings", () => { - it("keeps well-formed models and token limits and trims the workspace URL", () => { - const values: Record = { +describe("readConfig", () => { + it("keeps well-formed models, names and token limits and trims the workspace URL", () => { + const config = { workspace: "https://ws/", authCommand: "ug auth-token", models: [ @@ -98,20 +91,27 @@ describe("readSettings", () => { { id: "b", api: "messages" }, { api: "responses" }, null, - { id: "c", api: "chat_completions", maxInputTokens: 936_000, maxOutputTokens: "32000", extra: true }, - { id: "d", api: "chat_completions", maxInputTokens: 0, maxOutputTokens: 1.5 }, + { 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 }, ], }; - assert.deepEqual(readSettings((key) => values[key]), { + assert.deepEqual(readConfig(JSON.stringify(config)), { workspace: "https://ws", authCommand: "ug auth-token", models: [ { id: "a", api: "responses" }, - { id: "c", api: "chat_completions", maxInputTokens: 936_000 }, + { id: "c", name: "Model C", api: "chat_completions", maxInputTokens: 936_000 }, { id: "d", api: "chat_completions" }, ], }); }); + + it("gives no models for a missing or malformed file", () => { + const none = { workspace: "", authCommand: "", models: [] }; + for (const text of [undefined, "", "{", "[]", "null", '"text"']) { + assert.deepEqual(readConfig(text), none, String(text)); + } + }); }); describe("DatabricksModels", () => { @@ -119,19 +119,19 @@ describe("DatabricksModels", () => { workspace: "https://ws", authCommand: command, models: [ - { id: "system.ai.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: "chat_completions", maxInputTokens: 872_000, maxOutputTokens: 128_000 }, { id: "system.ai.gpt-6-luna", api: "responses" }, ], }); - it("lists the configured models, or none with a hint when settings are missing", async () => { + it("lists the configured models by name, or by id without one, or none with a hint when settings are missing", async () => { const logs: string[] = []; const models = await new DatabricksModels(() => settings("cmd"), (line) => logs.push(line)).provideLanguageModelChatInformation(); assert.deepEqual( - models.map((model) => [model.id, model.family, model.api, model.maxInputTokens, model.maxOutputTokens]), + 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", "chat_completions", 872_000, 128_000], - ["system.ai.gpt-6-luna", "gpt-6-luna", "responses", 128_000, 32_000], + ["system.ai.claude-sonnet-5-5", "Claude Sonnet 5.5", "claude-sonnet-5-5", "chat_completions", 872_000, 128_000], + ["system.ai.gpt-6-luna", "system.ai.gpt-6-luna", "gpt-6-luna", "responses", 128_000, 32_000], ], ); assert.deepEqual(await new DatabricksModels(() => ({ ...settings(""), models: [] }), (line) => logs.push(line)).provideLanguageModelChatInformation(), []); @@ -140,8 +140,8 @@ describe("DatabricksModels", () => { it("sends each model's turns to its API with the User-Agent", async () => { const requests = stubFetch( - [200, sse('{"choices":[{"delta":{"content":"hello"}}]}', "[DONE]")], - [200, sse('{"type":"response.output_text.delta","delta":"hi"}')], + [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(); @@ -156,10 +156,10 @@ describe("DatabricksModels", () => { ); } assert.deepEqual( - requests.map((request) => [request.url, request.body.model]), + requests.map((request) => [request.url, request.body.model, request.userAgent]), [ - ["https://ws/ai-gateway/mlflow/v1/chat/completions", "system.ai.claude-sonnet-5-5"], - ["https://ws/ai-gateway/mlflow/v1/responses", "system.ai.gpt-6-luna"], + ["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"], ], ); 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 9f66d8817..6abf2c81d 100644 --- a/vscode-extension/test/formats.test.ts +++ b/vscode-extension/test/formats.test.ts @@ -28,6 +28,14 @@ describe("Chat Completions", () => { }); }); + it("sends a tool's images in the user message after its result, since a tool message can't carry them", () => { + const request = chatCompletionsRequest({ ...turn, messages: [message(USER, result("c1", text("Screenshot:"), image(9)))] }); + assert.deepEqual((request as { messages: unknown[] }).messages, [ + { role: "tool", tool_call_id: "c1", content: "Screenshot:" }, + { role: "user", content: [{ type: "image_url", image_url: { url: "data:image/png;base64,CQ==" } }] }, + ]); + }); + it("gives a tool without an input schema an empty one, and omits tools when there are none", () => { const noSchema = chatCompletionsRequest({ ...turn, tools: [{ name: "list", description: "List", inputSchema: undefined }], toolRequired: false }); assert.deepEqual((noSchema as { tools: unknown[] }).tools, [ @@ -55,7 +63,7 @@ describe("Chat Completions", () => { chatCompletionsReplies( events( { choices: [{ delta: { content: [{ type: "reasoning", summary: [{ type: "summary_text", text: "The file may not exist." }] }] } }] }, - { choices: [{ delta: { content: [{ type: "text", text: "Checking." }, { type: "text", text: " Done." }] } }] }, + { choices: [{ delta: { content: [{ type: "text", text: "Checking." }, { type: "text", text: " Done." }] }, finish_reason: "stop" }] }, ), ), ); @@ -65,10 +73,25 @@ describe("Chat Completions", () => { it("fails on a length stop and on arguments that aren't an object", async () => { await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "length" }] }))), "output token limit"); await rejectsWith( - collect(chatCompletionsReplies(events({ choices: [{ delta: { tool_calls: [{ index: 0, id: "c", function: { name: "t", arguments: "[1]" } }] } }] }))), + collect( + chatCompletionsReplies( + events({ choices: [{ delta: { tool_calls: [{ index: 0, id: "c", function: { name: "t", arguments: "[1]" } }] }, finish_reason: "tool_calls" }] }), + ), + ), "aren't a JSON object", ); }); + + it("fails on an error event in the stream", async () => { + await rejectsWith(collect(chatCompletionsReplies(events({ error: { message: "overloaded" } }))), "overloaded"); + await rejectsWith(collect(chatCompletionsReplies(events({ error: "rate limited" }))), "rate limited"); + }); + + it("fails when the model stops for another reason or the stream ends without a finish reason", async () => { + await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "content_filter" }] }))), "stopped by a content filter"); + await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "insufficient_system_resource" }] }))), "(insufficient_system_resource)"); + await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ delta: { content: "Half an ans" } }] }))), "ended before the model finished"); + }); }); describe("estimateTokens", () => { @@ -98,6 +121,34 @@ describe("Responses", () => { }); }); + it("sends a tool's images inside its output", () => { + const request = responsesRequest({ ...turn, messages: [message(USER, result("c1", text("Screenshot:"), image(9)))] }); + assert.deepEqual((request as { input: unknown[] }).input, [ + { + type: "function_call_output", + call_id: "c1", + output: [ + { type: "input_text", text: "Screenshot:" }, + { type: "input_image", image_url: "data:image/png;base64,CQ==", detail: "auto" }, + ], + }, + ]); + }); + + it("holds function calls back until the response completes", async () => { + const seen: unknown[] = []; + const call = { type: "response.output_item.done", item: { type: "function_call", call_id: "c1", name: "get_weather", arguments: "{}" } }; + await rejectsWith( + (async () => { + for await (const reply of responsesReplies(events({ type: "response.output_text.delta", delta: "Checking." }, call))) { + seen.push(reply); + } + })(), + "ended before the model finished", + ); + assert.deepEqual(seen, [new LanguageModelTextPart("Checking.")]); + }); + it("streams text and each finished function call", async () => { const replies = await collect( responsesReplies( @@ -119,4 +170,11 @@ describe("Responses", () => { ); await rejectsWith(collect(responsesReplies(events({ type: "response.failed", response: { error: { message: "overloaded" } } }))), "overloaded"); }); + + it("fails on every other incomplete reason and on a stream without response.completed", async () => { + const incomplete = (reason?: string) => ({ type: "response.incomplete", response: { incomplete_details: reason ? { reason } : null } }); + await rejectsWith(collect(responsesReplies(events(incomplete("content_filter")))), "stopped by a content filter"); + await rejectsWith(collect(responsesReplies(events(incomplete()))), "(no reason given)"); + await rejectsWith(collect(responsesReplies(events({ type: "response.output_text.delta", delta: "Half an ans" }))), "ended before the model finished"); + }); });