[xi-db]

🥞 Stacked PR

Use this link to review incremental changes.

• stack/vscode-extension [Files changed] ← this PR
  • stack/vscode-configure [Files changed]

---

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.

[github-actions]

Polly AI Review

Blocking issues

• [P2] Force token renewal on the 401 retry — vscode-extension/src/gateway.ts:53. can the retry request a forced token refresh? rerunning the same helper can return the same cached, rejected OAuth token.

  invalidate() clears only extension memory; vscode-extension/src/auth.ts:79 reruns the unchanged command. The documented helper omits --force-refresh. For a server-rejected token whose exp is still in the future, the external cache can therefore defeat recovery. The existing TokenCache.refresh() in src/ucode/gateway_proxy.py:189 explicitly forces renewal for this case.

  Extend vscode-extension/test/gateway.test.ts:61, fetches a new token and retries once after a 401, with a helper that returns its cached token until explicitly refreshed.

Security vulnerabilities

• [P1] Repository settings can redirect the user’s bearer — vscode-extension/package.json:35. can we make the URL and auth helper machine-only? a trusted repository can override just the URL and receive the user-level helper’s bearer on the next request.

  Both unityGateway.workspace and unityGateway.authCommand omit configuration scope. vscode-extension/src/extension.ts:15 reads their merged values independently; vscode-extension/src/gateway.ts:38 combines the helper’s token with the configured destination.

  A repository can set only unityGateway.workspace in .vscode/settings.json to an attacker-controlled HTTPS origin, leaving the legitimate user-level helper unchanged. Selecting a configured model then sends that helper’s token to the overridden origin. This requires a trusted workspace and a chat request; it is not an untrusted-workspace escape claim.

  Add a configuration-isolation regression using dummy credentials. The existing injected settings reader does not exercise VS Code’s configuration precedence.

Non-blocking notes

• vscode-extension/package.json:82: can the staged VSIX assets include OpenAI’s applicable license/notices? the build strips legal comments, while staging copies only the MIT license.
• vscode-extension/test/auth.test.ts:71, reports the command's failure without its output: can this exercise the shell boundary? the injected error is already sanitized.
• vscode-extension/test/provider.test.ts:120, keeps one gateway client until the workspace or auth command changes: can this also change the workspace? currently only the auth-command transition is asserted.
• vscode-extension/test/vscode-mock.ts:49: can a mid-stream cancellation test use cancel() and check listener disposal, or can the unused instrumentation be removed?

Approach

The protocol split and shared message-part helpers fit the goal. Reuse the existing distinction between ordinary token retrieval and forced renewal for 401 recovery. No materially simpler replacement for the two wire-format adapters was identified.

Summary

Diff size (computed): 32 files; +3635 / -0 lines (3635 changed text lines); 0 binary files.

Category	Files	Added	Deleted	Share of changed text lines
Tests and test support	8	687	0	18.9%
Documentation	2	37	0	1.0%
Dependencies and lockfiles	3	1952	0	53.7%
Implementation / other	19	959	0	26.4%

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.

Changes

The PR adds provider registration, settings and model metadata, authentication, Chat Completions and Responses adapters, streaming conversion, logging, build tooling, CI, tests, licensing documentation, dependencies, and committed distribution assets.

The size is broadly justified by introducing a separate extension and shipping its bundled dependencies. The staged manifest/license copies have a concrete packaging purpose and a rebuild-drift check. Beyond the unused cancellation instrumentation noted above, no substantial removable duplication was identified.

Tests

Test-by-test assessment

This is a static assessment. Tests, typecheck, build, wheel contents, and byte-identical rebuilding were not executed or independently verified. Existing auth/proxy tests were inspected for context, but they do not establish the extension’s configuration precedence or external-helper refresh behavior.

Environment limitation: OpenCode, Cursor, Hermes, Pi, and Antigravity CLIs were unavailable on PATH.

Test / case	Behavior protected	Layer	Needed?	Action / rationale
vscode-extension/test/auth.test.ts:17 — reads exp from a JWT	JWT expiry conversion	Unit	keep	Small deterministic check, distinct from cache lifecycle behavior.
vscode-extension/test/auth.test.ts:21 — assumes an hour for a token without exp	Opaque-token fallback lifetime	Unit	keep	Covers the separate non-JWT credential path.
vscode-extension/test/auth.test.ts:27 — runs the command once and reuses the token until it nears expiry	Cache reuse and refresh buffer	Unit	keep	Injected time checks both sides of the freshness boundary without process setup.
vscode-extension/test/auth.test.ts:42 — fetches once for concurrent requests	Single-flight token fetching	Unit	keep	Concurrent behavior is not covered by sequential cache reuse.
vscode-extension/test/auth.test.ts:51 — fetches a new token after the current one is rejected	Extension-cache invalidation	Unit	keep	Protects the local invariant; does not establish renewal of the external helper’s cache.
vscode-extension/test/auth.test.ts:61 — ignores a rejection of a token it already replaced	Ignoring mismatched invalidation	Unit	keep	Tests rejection of a different token; the fixture does not actually fetch and replace that rejected token.
vscode-extension/test/auth.test.ts:71 — reports the command's failure without its output	Auth-error wrapping	Unit	keep	Wrapping coverage is useful, but the sanitized injected error does not prove shell-output redaction; see the note above.
vscode-extension/test/auth.test.ts:79 — rejects an empty token	Empty credential rejection	Unit	keep	Checks a meaningful failure before an authenticated request is attempted.
vscode-extension/test/auth.test.ts:85 — tells the caller when it fetched a token, with its expiry	Fetch-notification metadata	Unit	keep	Protects the logging callback contract without coupling to VS Code output.
vscode-extension/test/chat.test.ts:24 — converts a tool-using conversation with images	Chat history, tools, results, images, and option filtering	Unit	keep	One coherent conversation checks distinct wire transformations with proportionate setup.
vscode-extension/test/chat.test.ts:49 — gives a tool without an input schema an empty one	Chat tool-schema fallback	Unit	keep	Protects Chat Completions’ nested tool shape, unlike the Responses case.
vscode-extension/test/chat.test.ts:58 — sends no tools or tool choice when none are offered	Tool-free request construction	Unit	keep	Checks omission behavior, including a required-tool option without offered tools.
vscode-extension/test/chat.test.ts:67 — streams text and assembles tool calls split across chunks	Text forwarding and fragmented tool arguments	Unit	keep	Meaningful streaming assertions at the smallest reliable layer.
vscode-extension/test/chat.test.ts:85 — flags invalid tool arguments and a length stop	Malformed arguments and truncation markers	Unit	keep	Explicitly checks both adapter outputs; provider error presentation is separate.
vscode-extension/test/gateway.test.ts:40 — calls the workspace's mlflow routes with the token and ug's User-Agent	Routes, bearer headers, attribution, and token reuse	SDK component with fake fetch	keep	Exercises actual SDK request construction for both APIs without network access.
vscode-extension/test/gateway.test.ts:61 — fetches a new token and retries once after a 401	Auth retry and bearer replacement	SDK component with fake fetch	keep	Extend as requested above: sequential fake tokens currently mask the external-cache failure.
vscode-extension/test/gateway.test.ts:71 — asks the user to sign in when the new token is rejected too	Bounded retry and final guidance	SDK component with fake fetch	keep	Request-count assertions guard against retry loops; guidance is observable behavior.
vscode-extension/test/gateway.test.ts:81 — doesn't retry other errors	Non-auth error propagation	SDK component with fake fetch	keep	Protects the distinct policy against replaying unrelated failures.
vscode-extension/test/gateway.test.ts:91 — names ug's extension, Copilot Chat's version and the platform	User-Agent format and absent-version fallback	Unit	keep	Platform-tolerant assertions avoid an unnecessary OS matrix.
vscode-extension/test/provider.test.ts:72 — lists the configured models with their API	Provider model-information propagation	Component with VS Code mock	keep	Verifies picker-facing wiring beyond pure descriptor construction.
vscode-extension/test/provider.test.ts:83 — lists nothing and says why when settings are missing	Empty model list and configuration guidance	Component with VS Code mock	keep	Protects observable provider behavior beyond missing-setting detection.
vscode-extension/test/provider.test.ts:90 — calls Chat Completions or Responses by the model's API	API dispatch and VS Code response-part construction	Component with VS Code mock	keep	Both routes and tool-call conversion add coverage beyond payload tests.
vscode-extension/test/provider.test.ts:112 — turns a token-limit stop into an error	Truncation presentation	Component with VS Code mock	keep	Checks that adapter markers become errors rather than invalid response parts.
vscode-extension/test/provider.test.ts:120 — keeps one gateway client until the workspace or auth command changes	Client reuse and configuration transitions	Component with VS Code mock	keep	Auth-command transition is covered; add the workspace transition to this scenario as noted above.
vscode-extension/test/provider.test.ts:132 — announces model changes on refresh	Model-information change events	Component with VS Code mock	keep	Distinct lifecycle behavior; no live backend is needed.
vscode-extension/test/responses.test.ts:9 — sends the whole conversation statelessly, with tools and images	History replay, store: false, tools, results, and images	Unit	keep	Essential stateless-contract coverage for a different wire format.
vscode-extension/test/responses.test.ts:38 — gives a tool without an input schema an empty one	Responses tool-schema fallback	Unit	keep	Protects top-level function parameters rather than Chat Completions’ nested shape.
vscode-extension/test/responses.test.ts:52 — streams text and emits each finished function call	Responses event selection and completed calls	Unit	keep	Covers ignored events and non-function items alongside emitted results.
vscode-extension/test/responses.test.ts:73 — reports a token-limit stop	Responses truncation translation	Unit	keep	Protects a protocol-specific event shape.
vscode-extension/test/responses.test.ts:81 — throws on a failed response and on an error event	Both Responses failure-event forms	Unit	keep	Related failure paths have independent assertions; combining them is reasonable.
vscode-extension/test/settings.test.ts:7 — reads the workspace, models and auth command	Parsing and URL normalization	Unit	keep	Useful parsing coverage, but the injected reader cannot establish configuration isolation.
vscode-extension/test/settings.test.ts:29 — drops malformed model entries	Optional malformed-entry tolerance	Unit	keep	Mixed inputs check graceful recovery at the lowest reliable layer.
vscode-extension/test/settings.test.ts:35 — names every missing setting	Complete missing-setting diagnostics	Unit	keep	Pure contract distinct from provider presentation.
vscode-extension/test/settings.test.ts:41 — shows system.ai models by their short name	Prefix handling and Claude metadata	Unit	keep	Checks descriptor rules not fully asserted by provider propagation tests.
vscode-extension/test/settings.test.ts:50 — keeps other model names whole	Custom-name preservation	Unit	keep	Guards against stripping arbitrary qualified names.
vscode-extension/test/settings.test.ts:54 — treats gpt-oss as an open model without image input	GPT versus GPT-OSS capabilities	Unit	keep	Compact contrast protects prefix-ordering behavior.
vscode-extension/test/helpers.ts:1 — support: message/part builders, asStream, collect, jwt	Explicit fixtures and stream mechanics	Test support	keep	Reused across compatible scenarios; faults and assertions remain in tests.
vscode-extension/test/vscode-mock.ts:1 — support: response parts, constants, EventEmitter, token subscriptions	Minimal VS Code runtime seam	Test support	keep	Enables provider component tests without an extension host; does not establish real-host behavior.
vscode-extension/test/vscode-mock.ts:49 — support: unused cancellation/disposal instrumentation	Intended cancellation assertions	Test support	consolidate	Use it in the cancellation scenario requested above, or trim it; it currently establishes no coverage.

The two demonstrated defects need configuration-isolation and forced-refresh regressions. Existing adapter and SDK-component tests are otherwise appropriately layered; duplicating their assertions in live E2E tests is unnecessary. A focused extension-host check uniquely establishes real VS Code configuration precedence.

After fixes, run cd vscode-extension && npm run typecheck && npm test && npm run build, then check that rebuilding leaves the staged assets unchanged. For manual verification, use an isolated VS Code profile and dummy credentials to confirm repository settings cannot redirect the credential destination/helper; then check Ask and Agent tool/result turns through both API modes.

Scope

All inspected changes support the stated extension outcome; no clearly unrelated changes were found. Deferring ug configure integration is consistent with the described stack. The truncated description limits additional context, but does not establish unrelated work. The remaining packaging question is whether the independently distributed VSIX will carry the applicable bundled-dependency notices.

---

_{Automated review by Polly · workflow run}

[xi-db]

Thanks. Addressed in 9e8bdc9, which also rewrites the extension smaller: no runtime dependencies, four source files, a 10 KB bundle.

• [P1] Settings redirecting the token: valid. unityGateway.workspace, unityGateway.models and unityGateway.authCommand are now "scope": "application", so VS Code reads them only from user settings; a test checks the manifest keeps that. Checked in real VS Code 1.143.1: a workspace .vscode/settings.json pointing unityGateway.workspace at another host and unityGateway.authCommand at a different command was ignored. The chat turn completed through the user-level settings, and the workspace's command never ran.
• [P2] Forced renewal on the 401 retry: valid. After a 401 the extension now runs the auth command with --force-refresh (Tokens.renew), like TokenCache.refresh(). Tested against a real command that keeps printing its cached token until it's run with --force-refresh.
• Non-blocking: the bundled-dependency license note and the auth, provider and mock tests it names are gone with the rewrite: there's no openai SDK anymore, and the auth tests now run a real shell command. There's no mid-stream cancellation test; cancellation aborts the fetch through an AbortController tied to the request's cancellation token.

[masonc08]

approving to put in a working version but followups should be completed