Repository navigation
Conversation
| def _run_code(code: Path, *args: str): | ||
| return subprocess_cross_os.run( | ||
| [str(code), *args], capture_output=True, text=True, timeout=_CODE_TIMEOUT_S, check=False | ||
| ) |
There was a problem hiding this comment.
what happens if code hangs or can't be executed? TimeoutExpired/OSError aren't caught, so ug revert crashes before restore_file/clear_state runs. Configure only catches RuntimeError, so it crashes too. can we catch these and warn instead?
Suggestion:
| def _run_code(code: Path, *args: str): | |
| return subprocess_cross_os.run( | |
| [str(code), *args], capture_output=True, text=True, timeout=_CODE_TIMEOUT_S, check=False | |
| ) | |
| 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} {' '.join(args)}` failed: {exc}") from exc |
There was a problem hiding this comment.
Fixed in 918f4fa. _run_code now turns TimeoutExpired and OSError into the same RuntimeError as any other code failure, so configure warns instead of crashing. ug revert also wraps the VS Code step, so a failure there shows as failed and the other restores still run.
| if not models: | ||
| print_warning(f"No models for VS Code: {reason}") |
There was a problem hiding this comment.
Fixed in 918f4fa. When the listing comes back empty, configure warns with the reason and leaves VS Code unchanged (settings and extension). A working setup from an earlier run stays as it was, and a workspace switch can't pair the old models with the new workspace.
| workspace = state["workspace"] | ||
| models, reason = model_settings(workspace, launch_token(state, workspace)) | ||
| if not models: | ||
| print_warning(f"No models for VS Code: {reason}") |
There was a problem hiding this comment.
Fixed in 918f4fa. Without a reason it now says No Claude or GPT models for VS Code (the workspace has none). VS Code was left unchanged.
| from ucode.agents import vscode | ||
|
|
||
| vscode_result = vscode.revert() |
There was a problem hiding this comment.
Fixed in 918f4fa. Revert now runs code only when it found ug's unityGateway.* settings, since configure writes them before it installs. A failure in this step no longer stops the rest of revert either.
| if wants_vscode: | ||
| from ucode.agents import vscode | ||
|
|
||
| vscode.configure(load_state()) |
There was a problem hiding this comment.
Fixed in 918f4fa. --dry-run now skips VS Code with a note: a dry run doesn't save the workspace that VS Code would be set up from.
| @@ -0,0 +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=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},x=(e,n,t,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let a of B(n))!H.call(e,a)&&a!==t&&f(e,a,{get:()=>n[a],enumerable:!(s=F(n,a))||s.enumerable});return e};var C=(e,n,t)=>(t=e!=null?U(W(e)):{},x(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),Y=e=>x(f({},"__esModule",{value:!0}),e);var ne={};K(ne,{DatabricksModels:()=>v,activate:()=>te,readSettings:()=>N,userAgent:()=>q});module.exports=Y(ne);var c=C(require("vscode"));var b="0.1.0";var T=require("node:child_process"),X=5*6e4,Q=60*6e4,V=e=>new Promise((n,t)=>{(0,T.exec)(e,{timeout:6e4,windowsHide:!0},(s,a,o)=>{s?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(o).trim().slice(0,500)||s.message}`)):n(String(a).trim())})}),Z=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()+Q},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}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(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:Z(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var i=C(require("vscode")),_=3,ee=1e3,y=e=>e instanceof i.LanguageModelTextPart?e.value:e instanceof i.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof i.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",S=e=>e instanceof i.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,R=e=>e.inputSchema??{type:"object",properties:{}},L=(e,n,t)=>{let s;try{s=t.trim()?JSON.parse(t):{}}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: ${t}`);return new i.LanguageModelToolCallPart(e,n,s)},P=()=>new Error("The model stopped at its output token limit before finishing the response."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},E=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let s of e.content){let a=s instanceof i.LanguageModelToolResultPart?s.content:[s];for(let o of a)n+=o instanceof i.LanguageModelToolCallPart?o.name.length+JSON.stringify(o.input).length:y(o).length,t+=S(o)===void 0?0:1}return Math.ceil(n/4)+t*ee},j=e=>({results:e.filter(n=>n instanceof i.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof i.LanguageModelToolCallPart),text:e.map(y).join(""),images:e.map(S).filter(n=>n!==void 0)});function D(e){let n=e.messages.flatMap(t=>{let{results:s,calls:a,text:o,images:u}=j(t.content);if(Number(t.role)===_)return[{role:"system",content:o}];if(t.role===i.LanguageModelChatMessageRole.Assistant){let r=a.map(g=>({id:g.callId,type:"function",function:{name:g.name,arguments:JSON.stringify(g.input)}}));return o||r.length?[{role:"assistant",content:o||null,...r.length?{tool_calls:r}:{}}]:[]}let l=[...o?[{type:"text",text:o}]:[],...u.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("")})),...l.length?[{role:"user",content:l}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:R(t)}}))}}async function*O(e){let n=[];for await(let t of e)for(let s of t.choices??[]){s.delta?.content&&(yield new i.LanguageModelTextPart(s.delta.content));for(let a of s.delta?.tool_calls??[]){let o=n[a.index??0]??={id:"",name:"",args:""};o.id=a.id||o.id,o.name=a.function?.name||o.name,o.args+=a.function?.arguments??""}if(s.finish_reason==="length")throw P()}for(let t of n.filter(Boolean))yield L(t.id,t.name,t.args)}function I(e){let n=e.messages.flatMap(t=>{let{results:s,calls:a,text:o,images:u}=j(t.content);if(Number(t.role)===_)return[{role:"system",content:o}];if(t.role===i.LanguageModelChatMessageRole.Assistant)return[...o?[{role:"assistant",content:o}]:[],...a.map(r=>({type:"function_call",call_id:r.callId,name:r.name,arguments:JSON.stringify(r.input)}))];let l=[...o?[{type:"input_text",text:o}]:[],...u.map(r=>({type:"input_image",image_url:r,detail:"auto"}))];return[...s.map(r=>({type:"function_call_output",call_id:r.callId,output:r.content.map(y).join("")})),...l.length?[{role:"user",content:l}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:R(t),strict:!1}))}}async function*G(e){for await(let n of e){let t=n;if(t.type==="response.output_text.delta"&&t.delta)yield new i.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")yield L(t.item.call_id??"",t.item.name??"",t.item.arguments??"");else{if(t.type==="response.incomplete"&&t.response?.incomplete_details?.reason==="max_output_tokens")throw P();if(t.type==="response.failed"||t.type==="error")throw new Error(t.response?.error?.message??t.message??"The model response failed.")}}}async function*k(e,n,t,s,a,o){let u=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...a,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:o}),l=await s.get(),r=await u(l);if(r.status===401&&(await r.body?.cancel(),l=await s.renew(l),r=await u(l)),r.status===401)throw await r.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${e}\`, then retry.`);if(!r.ok||!r.body)throw new Error(`AI Gateway returned ${r.status}: ${(await r.text()).slice(0,300)}`);let g=new TextDecoder,p="";for await(let m of r.body){p+=g.decode(m,{stream:!0});let M=p.split(` | |||
| `);p=M.pop()??"";for(let J of M){let w=$(J);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let d=$(p+g.decode());d&&d!=="[DONE]"&&(yield JSON.parse(d))}var $=e=>e.startsWith("data:")?e.slice(5).trim():"";var N=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")),authCommand:String(e("authCommand")??"")}},q=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(n,t,s=()=>({})){this.settings=n;this.log=t;this.headers=s}settings;log;headers;changed=new c.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:s}=this.settings();return!n||!s||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:a,api:o})=>({id:a,api:o,name:a,family:a.split(".").pop()??a,version:b,detail:"Databricks",maxInputTokens:128e3,maxOutputTokens:32e3,capabilities:{toolCalling:!0,imageInput:o==="responses"||a.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,s,a,o){let{workspace:u,authCommand:l}=this.settings();this.tokens??=new h(l,d=>this.log(`Fetched a Databricks token (expires ${new Date(d).toISOString()})`));let r={model:n.id,messages:t,tools:s.tools??[],toolRequired:s.toolMode===c.LanguageModelChatToolMode.Required},g=new AbortController,p=o.onCancellationRequested(()=>g.abort());try{let d=n.api==="responses"?G(k(u,"responses",I(r),this.tokens,this.headers(),g.signal)):O(k(u,"chat/completions",D(r),this.tokens,this.headers(),g.signal));for await(let m of d)a.report(m)}catch(d){if(!o.isCancellationRequested)throw this.log(`${n.id} failed: ${d instanceof Error?d.message:String(d)}`),d}finally{p.dispose()}}async provideTokenCount(n,t){return E(t)}};function te(e){let n=c.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let o=c.workspace.getConfiguration("unityGateway");return N(u=>o.get(u))},s=()=>({"User-Agent":q(c.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),a=new v(t,o=>n.info(o),s);e.subscriptions.push(n,c.lm.registerLanguageModelChatProvider("databricks-unity-gateway",a),c.commands.registerCommand("unityGateway.refreshModels",()=>a.refresh()),c.workspace.onDidChangeConfiguration(o=>{o.affectsConfiguration("unityGateway")&&a.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); | |||
| "directory": "vscode-extension" | ||
| }, | ||
| "engines": { | ||
| "vscode": "^1.128.0" |
There was a problem hiding this comment.
Keeping ^1.128.0. 1.128 is the first VS Code that offers extension models in the Copilot harness as well as the Local one, and this extension relies on that. The types target 1.125 only because @types/vscode 1.128 wasn't published. On an older VS Code, the install fails with code's message as a warning, and the rest of configure still succeeds.
| ] | ||
| }, | ||
| "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 \"require('fs').copyFileSync('package.json', '../src/ucode/vscode/extension/package.json')\"", |
There was a problem hiding this comment.
Fixed in 1d76882. The build now writes the shipped package.json without scripts and devDependencies.
| def test_finds_the_macos_app_when_code_isnt_on_path(self, monkeypatch, tmp_path): | ||
| monkeypatch.setattr(vscode.shutil, "which", lambda name: None) | ||
| monkeypatch.setattr(vscode, "current_os", lambda: OS.MACOS) | ||
| monkeypatch.setattr(vscode.Path, "home", lambda: tmp_path) |
There was a problem hiding this comment.
Fixed in 918f4fa. The lookup is now tested in two parts, and neither reads the real /Applications: one checks the locations it looks in, in order, and the other checks the fallback to the first one that exists, using stubbed locations.
| monkeypatch.setattr( | ||
| "ucode.cli._configure_shared_workspace_states", | ||
| lambda *args, **kwargs: shared.append(kwargs) or [], | ||
| ) |
There was a problem hiding this comment.
Keeping it, and it now covers --dry-run too (918f4fa). Checking the outcome of the real path would need a signed-in workspace and an installed VS Code. Writing the settings and installing the extension are covered in TestConfigure. This test pins the CLI routing: --agents vscode sets up the workspace without the model-agent path, then hands that state to VS Code.
Polly AI ReviewBlocking issues
Security vulnerabilitiesNo verified security vulnerability found. The command and destination settings are application-scoped, preventing ordinary repository workspace settings from overriding them. Live security and dependency-advisory checks were not performed. Non-blocking notes
ApproachThe module boundaries and JSONC-preserving editor fit the existing patterns. No existing JSONC editor was found to reuse. For cleanup, follow the existing Desktop failure-result pattern so optional integration failures do not interrupt unrelated restoration. SummaryDiff size (computed): 27 files; +2489 / -7 lines (2496 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe PR connects workspace configuration to a bundled VS Code Copilot provider. The full snapshot includes the prerequisite extension, Python installation and cleanup, JSONC editing, documentation, test support, build checks, dependency metadata, a lockfile and generated JavaScript. The size is largely explained by delivering both the extension and its configure integration in this stacked snapshot. No concrete redundant implementation or scaffolding was found that could be removed while preserving the intended behavior. TestsTest-by-test assessmentAssessment covers all eight changed test/support files. Existing CLI cleanup-failure coverage provides a suitable place to extend restoration checks. The demonstrated serializer and stream failures belong in the existing unit/component tests, not additional E2Es. Three extension failures were reproduced through isolated synthetic invocations of the bundled extension, with network and process execution stubbed and credentials unavailable. Repository test suites, builds and live VS Code flows were not run; Python findings are based on control-flow inspection.
Real installation, activation and Copilot model selection remain unverified host boundaries. To check them manually, run ScopeNo clearly unrelated changes were found. The prerequisite extension, generated bundle, dependencies, configure integration, tests and documentation support the stated Copilot Chat setup outcome. The truncated description limits assessment of its reported validation, but does not establish a scope violation. Automated review by Polly · workflow run |
5dfb9b9 to
b756487
Compare
…ilot Chat (databricks#1314) ## 🥞 Stacked PR Use this [link](https://github.com/databricks/unity-gateway/pull/1314/files) to review incremental changes. - [**stack/vscode-extension**](databricks#1314) [[Files changed](https://github.com/databricks/unity-gateway/pull/1314/files)] ← _this PR_ - [stack/vscode-configure](databricks#1316) [[Files changed](https://github.com/databricks/unity-gateway/pull/1316/files/238af81c0cb95b41cbafb910d00da82005f181c5..3904468a7bfada1c978a78bf69ffe2e3d29d98ce)] --------- ## What First of two PRs that bring Databricks models to Copilot Chat in VS Code. This PR adds the extension; the next one makes `ug configure --agents vscode` pack, install and configure it. The extension (`vscode-extension/`, TypeScript, no runtime dependencies) registers a language model chat provider. It puts the workspace's AI Gateway models in the chat model picker, in both Ask and Agent mode and in both the Local and Copilot harnesses. It calls the gateway directly with `fetch`, with no local proxy: - **Chat Completions** (`{workspace}/ai-gateway/mlflow/v1/chat/completions`) for Claude. The next PR's `ug configure` offers only Claude and GPT models, the same set as Copilot CLI. - **Responses** (`{workspace}/ai-gateway/mlflow/v1/responses`) for GPT-5 and later. It's stateless: `store: false`, and the whole conversation is sent each turn. GPT-5.6 and later reject function tools on Chat Completions, as do the codex and pro models, so they need this route for Agent mode. - Streaming, tool calls, tool results and images work on both. - While a Claude model thinks, the gateway sends the content of a Chat Completions chunk as blocks rather than a string. The extension takes the text out of them and drops the reasoning: handing a block to Copilot Chat fails the turn. Settings, which `ug configure` will write to the user settings: | Setting | Value | |---|---| | `unityGateway.workspace` | Workspace URL | | `unityGateway.models` | `[{"id": "<model service>", "api": "chat_completions" \| "responses", "maxInputTokens"?: n, "maxOutputTokens"?: n}]` | | `unityGateway.authCommand` | The `ug auth-token --host ...` command ug already writes for other agents | All three are `"scope": "application"`, so VS Code reads them from user settings only. A repository's `.vscode/settings.json` can't redirect the token to another URL, or pick the command that prints it. **Auth.** The extension runs the auth command when it needs a token, keeps the token in memory until five minutes before its `exp`, and shares one fetch between concurrent requests. When the gateway answers 401 it runs the command again with `--force-refresh`, because the Databricks CLI can return its cached token while `exp` is still ahead, and retries once. If that also fails, it reports "Databricks rejected the sign-in. Run `databricks auth login --host ...`, then retry." It stores no credentials, and the token is never logged. **Other behavior:** - **User-Agent:** `ucode/<extension version> githubcopilotchat/<Copilot Chat version>`. The gateway already attributes the `githubcopilotchat` token to Copilot Chat. - **Model info:** the model `family` is the model's own name (`claude-sonnet-5-5`, `gpt-6-luna`), which Copilot Chat uses to pick its prompts and tools. The gateway doesn't report context limits, so a model's limits come from its `maxInputTokens` / `maxOutputTokens` settings, and a model without them gets 128K input and 32K output (Copilot's own default). Image input is on for Claude and Responses models. Token counts are about four characters per token, plus a flat 1,000 per image. - **Icon:** the provider shows the Databricks logo (`media/databricks.svg`, in the brand red) in the model picker. - **`Unity Gateway: Refresh Models`** re-reads the settings; settings changes also refresh. - **Logging:** the **Unity Gateway** output channel logs token fetches (expiry time only) and failed requests. - **`extensionKind: ["ui"]`**, engine `^1.128.0` (the first version that offers extension models in the Copilot harness). **Packaging.** - `npm run build` (esbuild, minified) writes `src/ucode/vscode/extension/out/extension.js` (9 KB) and copies `package.json` and `media/` next to it. Both are committed, so the wheel ships them and ug can pack the vsix with the standard library: no Node on users' machines, no binary in the repo. - `just vscode-build` rebuilds them, and rebuilds are byte-identical. - `.gitattributes` marks the bundle as generated. **CI.** A new `VS Code extension` job runs `npm ci` (with the npm cache), typecheck and the tests. It then rebuilds and fails if the committed files under `src/ucode/vscode/extension/` changed. The `test` check that branch rules require now needs this job too. `package-lock.json` records the public npm registry; `.npmrc` sets `replace-registry-host=always`, so `npm ci` fetches the same tarballs from whatever registry is configured locally. ## Not in this PR - `ug configure` doesn't install or configure the extension yet; that's the next PR. - No MCP registration for VS Code. - Claude goes through Chat Completions, so Anthropic-only features, prompt caching among them, aren't used. ## Validation - `npm run typecheck` is clean. `npm test` runs 25 cases with Node's own test runner (esbuild compiles them, with `vscode` aliased to a test stand-in, so the only dev dependencies are esbuild, TypeScript and type packages): - **auth**, against a real shell command: reuse until near expiry, one run for concurrent requests, `--force-refresh` after a rejection but not after a stale one or while a fetch is in flight, command failure and empty output; - **both request formats**: message, tool, tool-result and image conversion; tools without an input schema; fragment assembly; a thinking model's content blocks; token-limit stops; arguments that aren't an object; failed responses; token estimates that don't count image bytes; - **the gateway client**: SSE across chunk boundaries, `[DONE]` and a last event without a trailing newline; one forced retry after 401, then the sign-in message; other errors with the gateway's body; - **the provider**: model list with configured or default token limits, settings with malformed limits, missing settings, routing by `api` with the User-Agent; - **the manifest and User-Agent**: every setting is application-scoped; the provider's icon file exists; the User-Agent names the extension and Copilot Chat versions. - `just lint` passes. The PR changes no Python. - **By hand** with real VS Code 1.143.1 (Copilot Chat 0.71.0), signed out, under xvfb on Linux, against a real workspace. The extension was installed from a vsix packed from the committed files, and the settings were written to user settings by hand. Each case asked Agent mode to create a file with a given value: - **Claude** (`claude-haiku-5-5`, Chat Completions): file written, in both the Local and the Copilot harness. - **GPT** (`gpt-6-luna`, Responses): file written, in both the Local and the Copilot harness. - **`deepseek-v4-1-flash`** (Chat Completions): file written. - **Workspace override:** a workspace `.vscode/settings.json` set `unityGateway.workspace` to another host and `unityGateway.authCommand` to a different command. Both were ignored: the task completed through the user-level settings, and the workspace's command never ran. - **`gpt-oss-120b`**, tested on an earlier version: with Copilot's full 49-tool request, the model reasons about the tool call and then stops without making it. The gateway's response contains no tool call. - Not run: Windows or macOS, or a session longer than one token lifetime. The over-an-hour check comes with the configure PR. This pull request and its description were written by Isaac. --------- Co-authored-by: Isaac <no-reply@databricks.com>
`ug configure --agents vscode` configures the workspace the way Cursor's path does (no model picker), then: - finds VS Code's `code` command (or `code-insiders`), on PATH or where the installer puts it, since macOS doesn't add it to PATH by default; - packs the extension bundled with ug into a vsix with the standard library and installs it, unless that version is already installed; - writes unityGateway.workspace, unityGateway.models and unityGateway.authCommand into the user settings.json. Models come from discover_model_services: Claude, GPT and open models (not Gemini), GPT-5 and later on the Responses API, the rest on Chat Completions. The auth command is the one Claude Code already uses. The settings are edited with a new JSONC helper in config_io that changes only ug's top-level keys and keeps the user's comments and formatting. `ug revert` removes those keys and uninstalls the extension. Co-authored-by: Isaac <no-reply@databricks.com>
The extension now ships its provider icon next to the bundle and its manifest. Pack every file under the bundled extension folder instead of naming the two, so a new asset doesn't need a change here. Co-authored-by: Isaac <no-reply@databricks.com>
Configure skipped the install when VS Code already had the same extension version. A ug upgrade can change the bundled extension without changing its version, and users then kept the old build. Install with --force every time. Co-authored-by: Isaac <no-reply@databricks.com>
Each Claude entry in unityGateway.models now carries maxInputTokens and maxOutputTokens from the limits ug already gives Copilot CLI (copilot_model_limits), e.g. 936K / 32K for Opus 5.5 and 872K / 128K for Sonnet 5.5. Those limits hold on the Chat Completions route the extension uses: the gateway took an 861K-token prompt to claude-haiku-5-5 and a 128K max_tokens on Sonnet 5.5. GPT and open models get no entry, so they keep the extension's 128K / 32K. Co-authored-by: Isaac <no-reply@databricks.com>
VS Code applies the settings right away, but a window that's already open keeps running the extension build it loaded, and configure now reinstalls on every run. Print a note to reload open windows. Co-authored-by: Isaac <no-reply@databricks.com>
"Reload open VS Code windows (Developer: Reload Window)" assumed people know
the Command Palette. Print the steps instead, with this OS's shortcuts:
Next, in VS Code:
1. If VS Code is already open, reload each window:
press Cmd+Shift+P, type "Reload Window", and press Enter.
Or quit VS Code (Cmd+Q) and open it again.
2. Open Chat, click the model picker, and choose a Databricks model.
The reload matters only for a window that's already open: in VS Code
1.143, a first install shows up in an open window within seconds, but a
reinstall doesn't, and the window keeps running the build it loaded.
Co-authored-by: Isaac <no-reply@databricks.com>
Reloading each window through the Command Palette was more to explain, and
it misses any window the user doesn't reload. Restarting VS Code loads the
new extension build in every window at once:
Next, in VS Code:
1. If VS Code is open, quit it (Cmd+Q) and open it again.
2. Open Chat, click the model picker, and choose a Databricks model.
Windows and Linux say "close every VS Code window and open it again".
Co-authored-by: Isaac <no-reply@databricks.com>
Configure wrote every open model the workspace lists, but many of them don't work from VS Code. On e2-dogfood, five that the listing returns (deepseek-v4-flash-0731, deepseek-v4-pro-0813, glm-5-3, glm-5-3-flash, kimi-k3) fail with "Could not route request" on both Chat Completions and Responses, kimi-k3-neo returns 404, and gpt-oss-120b stops without calling a tool when Copilot sends its full tool list. Keep the models Copilot CLI offers: Claude and GPT, using the same rule (copilot_exclusion_reason), which leaves out gpt-oss and the other open models. Co-authored-by: Isaac <no-reply@databricks.com>
3904468 to
8b22c4c
Compare
Both stream readers stopped when the connection closed, so a stream cut off mid-answer (a dropped connection, a proxy timeout) reached Copilot Chat as a finished reply. Only an output-token stop was reported: a content-filter or any other early stop was ignored, and the user got a truncated answer with no reason. A reply now counts as done only after the gateway says the model finished: response.completed on Responses, finish_reason "stop" or "tool_calls" on Chat Completions. Any other stop reason fails with a message for it, and a stream that ends without a final event fails as cut off. Co-authored-by: Isaac <no-reply@databricks.com>
The picker listed models by their service ids, e.g.
system.ai.claude-sonnet-5-5. ug configure now writes a name for each
model with model_picker_label, the helper Codex's picker already uses
("Claude Sonnet 5.5", "GPT 6.1 Sol"), and the extension shows it. As in
Codex, two models with the same label keep their full ids. The id still
goes to the gateway, and settings without a name show the id as before.
Co-authored-by: Isaac <no-reply@databricks.com>
The build copied vscode-extension/media/databricks.svg next to the bundle, so the repo held two identical copies. The SVG has no build step and only the shipped extension uses it, so keep it only in src/ucode/vscode/extension/media/ and stop copying it. package.json still needs both copies (npm reads one, VS Code the other); mark the shipped copy as generated, like the bundle. Co-authored-by: Isaac <no-reply@databricks.com>
Review fixes for the ug side of VS Code: - A `code` command that hangs or can't start raised TimeoutExpired or OSError, which configure and revert didn't catch. It now fails like any other `code` error, with a warning. - A VS Code failure in `ug revert` no longer stops the Claude, Codex and Pi restores that run after it; it shows as "failed". - An empty model listing (a failed request, or a workspace without Claude or GPT models) wrote `unityGateway.models: []` over a working setup and printed "No models for VS Code: None". Configure now leaves VS Code unchanged and says why. - `ug revert` ran `code --list-extensions` for everyone. It now runs `code` only when ug's settings are there, since configure writes them before it installs. - `--dry-run` set VS Code up from the saved state, which a dry run doesn't update: a KeyError on a fresh machine, or the previous workspace. A dry run now leaves VS Code alone. - Settings files that can't be read or written warn instead of raising. - The macOS lookup test no longer depends on whether this Mac has VS Code in /Applications. Co-authored-by: Isaac <no-reply@databricks.com>
… extension
Review fixes for the extension:
- Images a tool returns (a screenshot, an image file it read) were dropped:
tool results were flattened to text. Responses now sends them inside the
function_call_output. On Chat Completions the gateway rejects images in a
tool message, so they go in the user message right after the results.
- An error event in a Chat Completions stream ({"error": ...}) was ignored;
it now fails the reply with the gateway's message.
- Responses tool calls are held back until response.completed, so a
cut-off stream can't start a tool, as on Chat Completions.
- After a second 401 the message said to run `databricks auth login`, which
is wrong for PAT and custom OAuth setups. It now points to
`ug configure --agents vscode`, which signs in whichever way ug is set up.
- The shipped package.json no longer carries the repo's build scripts and
dev dependencies.
- The provider test now checks the User-Agent each request sends.
Co-authored-by: Isaac <no-reply@databricks.com>
|
Thanks, Polly. Where each item stands (all fixed in this PR):
From the non-blocking notes: the provider test now asserts the User-Agent each request sends (1d76882). This comment was written by Isaac. |
🥞 Stacked PR
Use this link to review incremental changes.
What
Second of two PRs that bring Databricks models to Copilot Chat in VS Code. The first, #1314, adds the extension. This one makes
ug configure --agents vscodeset it up, so users never install anything in VS Code by hand.ug configure --agents vscode:Configures the workspace the way Cursor's path does: sign-in and workspace state, no model picker. Other agents can be configured in the same
--agentslist.Finds VS Code's
codecommand:code, thencode-insiders, on PATH or where VS Code's installer puts it. macOS doesn't putcodeon PATH unless the user runs "Shell Command: Install 'code' command in PATH". The install locations checked:/Applications/or~/Applications/Visual Studio Code.app/.../bin/code%LOCALAPPDATA%\Programs\Microsoft VS Code\bin\code.cmd, or the same under Program Files/usr/binor/snap/binIf VS Code isn't found, ug warns and the rest of configure still succeeds.
Writes the extension's settings into the user
settings.jsonfor that VS Code:unityGateway.workspace,unityGateway.modelsandunityGateway.authCommand.discover_model_services: Claude (newest per family) and GPT, filtered with Copilot CLI's rule (copilot_exclusion_reason), so VS Code offers the same models asug copilot. Gemini and open models,gpt-ossincluded, are left out (see Validation).responses, through the existing Copilot rule (now exposed ascopilot.uses_responses_api); everything else useschat_completions.copilot_model_limits): 936K in / 32K out for Opus 5.5 and Fable 5.1, 872K / 128K for Sonnet 5.5 and Haiku 5.5. VS Code then shows a 1M context window for them, and Copilot Chat doesn't compact a conversation at 128K. GPT models get none and keep the extension's 128K / 32K.build_auth_shell_command, or the custom OAuth form).Installs the extension. ug packs the bundled files from
src/ucode/vscode/extension/into a vsix with the standard library (no Node, novsce, about 2 ms). It then runscode --install-extension <vsix> --forceand deletes the temp file. It reinstalls on every run, even over the same version, because a ug upgrade can change the bundled extension without changing its version. Afterwards it prints what to do next:On Windows and Linux, step 1 says "close every VS Code window and open it again". A restart is needed only when VS Code was already open, and only after a reinstall. In VS Code 1.143, a first install showed up in an open window within about 3 seconds. After a reinstall, an open window kept serving the build it had loaded until it was reloaded or restarted (checked by installing a build with a different default token limit over the old one). Restarting covers every open window at once.
ug revertremoves ug's three settings and uninstalls the extension.JSONC editing. VS Code's
settings.jsonallows comments and trailing commas, and the repo had no helper for that.config_iogainsloads_jsoncandupdate_jsonc_keys. The editor tokenizes strings,//and/* */comments, and structure, then changes only the spans of the top-level keys it sets or removes. Everything else stays byte-for-byte, comments included:The unit-test isolation fixture now stubs out VS Code discovery and its settings path, so no test runs the developer's
codeor touches their settings.Follow-ups from the #1314 review
Four review comments on #1314, fixed here because #1314 had already merged:
response.completed(Responses) or afinish_reasonofstoportool_calls(Chat Completions). A stream that ends without one fails with "The response ended before the model finished. Try again."max_output_tokensandlengthused to throw, so acontent_filterstop silently truncated the answer. Every otherresponse.incompletereason andfinish_reasonnow fails with its own message: output token limit, content filter, or "stopped before finishing (<reason>)".ug configurewrites anamefor each model withmodel_picker_label, the helper Codex's picker uses ("Claude Sonnet 5.5", "GPT 6.1 Sol"), and the extension shows it. As in Codex, models whose labels collide keep their full ids. The id still goes to the gateway, and settings without anameshow the id.src/ucode/vscode/extension/media/, the folder ug ships, and the build no longer copies it. The shippedpackage.jsoncopy is marked generated, like the bundle.Fixes from the automated review
From github-actions and Polly (replied in each thread):
codefailures: acodecommand that hangs or can't start now warns instead of crashing configure or revert.ug revertshows asfailedand no longer stops the other restores.codewhen needed: it runscodeonly when ug's settings are there.--dry-run: VS Code is skipped, because a dry run doesn't save the workspace it would be set up from.function_call_output. Chat Completions sends them in the user message after the results, because the gateway rejects images in a tool message. In VS Code, Claude Haiku 5.5 and GPT 6 Luna each read a PNG throughcopilot_viewImageand named its colour.errorevent in a Chat Completions stream fails the reply with its message. Responses tool calls wait forresponse.completed.ug configure --agents vscode, notdatabricks auth login.package.json: it no longer includes the repo'sscriptsanddevDependencies.Not in this PR
Validation
tests/test_agent_vscode.py(18) andtests/test_config_io.py::TestJsonc(12) cover:code: on PATH, Insiders, the macOS app, the Windows user install, none;--agents vscodeCLI wiring;//, trailing commas, empty files, invalid files.test_agents_initupdate prompts,test_e2e_user_agent,test_smart_routertoggles) fail the same way onmain.just lintpasses.ug configure --agents vscode --profile <profile>foundcode, wrote 35 models (open models were still included then; the same workspace now gets 23) and the auth command into~/.config/Code/User/settings.json, and installeddatabricks.unity-gateway0.1.0 from the vsix it packed.claude-haiku-5-5(Chat Completions) andgpt-6-luna(Responses), with tokens from the configured auth command.ug revertthen reported the VS Code extension and settings as removed. The threeunityGateway.*keys were gone, the other 7 settings in the file were untouched, and the extension was uninstalled.claude-haiku-5-566 minutes apart each appended a line to a file. The first token expired at 18:58 UTC; the second turn, at 19:47, ran the auth command again for a new token, and both lines were written.DATABRICKS_BEARERto sign in, because the workspace-only path forcesdatabricks auth loginand the test box has no browser (Cursor's path already behaves this way). VS Code itself used theug auth-token --profilecommand.gpt-oss-120b,gpt-oss-20b,deepseek-v4-1-flash,deepseek-v4-1-flash-1017,glm-5-2andkimi-k2-7-codeworked.deepseek-v4-flash-0731,deepseek-v4-pro-0813,glm-5-3,glm-5-3-flashandkimi-k3failed with500 Could not route requeston both Chat Completions and Responses, with or without tools, andkimi-k3-neoreturned 404.gpt-oss-120balso stops without calling a tool when Copilot sends its full tool list. So VS Code offers what Copilot CLI offers.claude-haiku-5-5, 630K-token prompts to Haiku and Sonnet 5.5, andmax_tokensof 128K on Haiku and Sonnet 5.5 and 32K on Opus 5.5 and Fable 5.1. Near-1M prompts can still hit a workspace's per-minute input-token rate limit (a 429), which is a separate limit. In VS Code, Manage Language Models lists the four Claude models at 1M and the rest at 160K.npm testruns 27 cases, including streams cut off on each route andcontent_filterstops;tests/test_agent_vscode.pycovers the names and the label-collision fallback.tests/integration/journey. The integration runners have no VS Code, and launching VS Code is outside the integration scope intests/AGENTS.md. Also not run: Windows or macOS.This pull request and its description were written by Isaac.