copilot: add a VS Code extension that serves Databricks models to Copilot Chat

.gitattributes

@@ -0,0 +1 @@
+src/ucode/vscode/extension/out/extension.js linguist-generated=true

.github/workflows/ci.yml

@@ -89,6 +89,35 @@ jobs:
           export UCODE_TEST_CLAUDE_BINARY="$(command -v claude)"
           uv run pytest tests/native_claude/test_mod.py tests/native_claude/test_credentials.py -q
 
+  vscode-extension:
+    name: VS Code extension
+    runs-on: ubuntu-latest
+    needs: lint
+    timeout-minutes: 10
+    defaults:
+      run:
+        working-directory: vscode-extension
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          persist-credentials: false
+      - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020  # v4
+        with:
+          node-version: 22
+          cache: npm
+          cache-dependency-path: vscode-extension/package-lock.json
+      - run: npm ci --no-audit --no-fund
+      - run: npm run typecheck
+      - run: npm test
+      # ug packs the extension from these committed files, so they must match the sources.
+      - name: Check the committed bundle is up to date
+        run: |
+          npm run build
+          test -z "$(git status --porcelain -- ../src/ucode/vscode/extension)" || {
+            echo "::error::The VS Code extension bundle is stale. Run \`just vscode-build\` and commit the result."
+            exit 1
+          }
+
   e2e-shards:
     name: ${{ matrix.name }} · ${{ matrix.os == 'windows-latest' && 'Windows' || 'Linux' }}
     runs-on: ${{ matrix.os }}
@@ -200,14 +229,15 @@ jobs:
   required-tests:
     name: test
     if: ${{ !cancelled() }}
-    needs: test
+    needs: [test, vscode-extension]
     runs-on: ubuntu-latest
     timeout-minutes: 5
     steps:
-      - name: Require unit tests to pass
+      - name: Require unit and VS Code extension tests to pass
         env:
           RESULT: ${{ needs.test.result }}
-        run: test "$RESULT" = success
+          VSCODE_EXTENSION_RESULT: ${{ needs.vscode-extension.result }}
+        run: test "$RESULT" = success && test "$VSCODE_EXTENSION_RESULT" = success
 
   required-e2e:
     name: e2e

justfile

@@ -26,3 +26,7 @@ ty:
 fix:
     uv run ruff check --fix src/ tests/
     uv run ruff format src/ tests/
+
+# Rebuild the VS Code extension bundle ug packs at configure time. Commit the result.
+vscode-build:
+    cd vscode-extension && npm ci --no-audit --no-fund && npm run build

src/ucode/vscode/extension/package.json

@@ -0,0 +1,105 @@
+{
+  "name": "unity-gateway",
+  "displayName": "Databricks Unity Gateway",
+  "description": "Chat with Databricks AI Gateway models in VS Code through Unity Gateway (ug)",
+  "version": "0.1.0",
+  "publisher": "databricks",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/databricks/unity-gateway.git",
+    "directory": "vscode-extension"
+  },
+  "engines": {
+    "vscode": "^1.128.0"
+  },
+  "categories": [
+    "AI",
+    "Chat"
+  ],
+  "extensionKind": [
+    "ui"
+  ],
+  "main": "./out/extension.js",
+  "activationEvents": [],
+  "contributes": {
+    "languageModelChatProviders": [
+      {
+        "vendor": "databricks-unity-gateway",
+        "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"
+  }
+}

vscode-extension/.gitignore

@@ -0,0 +1,3 @@
+node_modules/
+*.vsix
+.test-out/

vscode-extension/.npmrc

@@ -0,0 +1,3 @@
+# The lockfile records the public registry; this lets `npm ci` fetch the same tarballs from
+# whichever registry is configured locally.
+replace-registry-host=always

vscode-extension/README.md

@@ -0,0 +1,22 @@
+# 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.
+
+`ug configure` installs it and writes its settings:
+
+| Setting | Value |
+|---|---|
+| `unityGateway.workspace` | The workspace URL |
+| `unityGateway.models` | `[{"id": "<model service>", "api": "chat_completions" \| "responses"}]` |
+| `unityGateway.authCommand` | A command that prints a Databricks access token (`ug auth-token --host ...`) |
+
+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.
+
+## Development
+
+`npm ci`, then `npm run typecheck`, `npm test` and `npm run build`. The build writes the bundle and
+`package.json` to `../src/ucode/vscode/extension/`, which ug packs from; commit them (`just vscode-build`).