diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 000000000..e21d1145a --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +src/ucode/vscode/extension/out/extension.js linguist-generated=true diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 24dbaab06..af34fb5d1 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/justfile b/justfile index f1fb8558f..faa6c7098 100644 --- a/justfile +++ b/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 diff --git a/src/ucode/vscode/extension/media/databricks.svg b/src/ucode/vscode/extension/media/databricks.svg new file mode 100644 index 000000000..b7bd91e12 --- /dev/null +++ b/src/ucode/vscode/extension/media/databricks.svg @@ -0,0 +1 @@ + diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js new file mode 100644 index 000000000..575c17279 --- /dev/null +++ b/src/ucode/vscode/extension/out/extension.js @@ -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=(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}); diff --git a/src/ucode/vscode/extension/package.json b/src/ucode/vscode/extension/package.json new file mode 100644 index 000000000..8bc404b46 --- /dev/null +++ b/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" + } +} diff --git a/vscode-extension/.gitignore b/vscode-extension/.gitignore new file mode 100644 index 000000000..13f124c2b --- /dev/null +++ b/vscode-extension/.gitignore @@ -0,0 +1,3 @@ +node_modules/ +*.vsix +.test-out/ diff --git a/vscode-extension/.npmrc b/vscode-extension/.npmrc new file mode 100644 index 000000000..8f1943004 --- /dev/null +++ b/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 diff --git a/vscode-extension/README.md b/vscode-extension/README.md new file mode 100644 index 000000000..d49c502a9 --- /dev/null +++ b/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": "", "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`). diff --git a/vscode-extension/media/databricks.svg b/vscode-extension/media/databricks.svg new file mode 100644 index 000000000..b7bd91e12 --- /dev/null +++ b/vscode-extension/media/databricks.svg @@ -0,0 +1 @@ + diff --git a/vscode-extension/package-lock.json b/vscode-extension/package-lock.json new file mode 100644 index 000000000..dfe196058 --- /dev/null +++ b/vscode-extension/package-lock.json @@ -0,0 +1,543 @@ +{ + "name": "unity-gateway", + "version": "0.1.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "unity-gateway", + "version": "0.1.0", + "devDependencies": { + "@types/node": "^22.20.3", + "@types/vscode": "1.125.0", + "esbuild": "^0.28.2", + "typescript": "^5.9.3" + }, + "engines": { + "vscode": "^1.128.0" + } + }, + "node_modules/@esbuild/aix-ppc64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.28.2.tgz", + "integrity": "sha512-XExcO+dvLKvVtNTibSTBej1NCAbaGhWn9Ww1ZPx80qsahhPFe/8jgWP0IchNe0F3HwkU7n8ejhH8bjonqht8mQ==", + "cpu": [ + "ppc64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "aix" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/android-arm": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/android-arm/-/android-arm-0.28.2.tgz", + "integrity": "sha512-kXXoiPVVGQcnIYGOeaovwOURpniDBpSq4A03qkQ+BMQqtGG6HYap3xne9C1O1yo4TR3qxlCX5IqqmX6fFo2Lqg==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/android-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/android-arm64/-/android-arm64-0.28.2.tgz", + "integrity": "sha512-5YfKeeI8qWfBZIX+u2xZC3Zlb3Os/gLS2sbEKM+I4ZOcsWmHS2WLysCcQZDAFRslDUU5Oiq44gf6PYN1vGwG5A==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/android-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/android-x64/-/android-x64-0.28.2.tgz", + "integrity": "sha512-O387ite7SzUyCcy3JQX4P4bLtEA7bLLkx+esve5JHnyYfNTxcVpXZo9jhdB0lTKN44gztELTdU7nS8Nr16Fs1Q==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/darwin-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/darwin-arm64/-/darwin-arm64-0.28.2.tgz", + "integrity": "sha512-n4KqkOQrraxHJcgjM1RvwbigfQKIKJVpM7xp+KsxiyUSrRdIXnt73VhrPAx0fV44hgfmIVKjxMN9J1t5jySVkw==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/darwin-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/darwin-x64/-/darwin-x64-0.28.2.tgz", + "integrity": "sha512-uq6suIWYP37qzGddBKPw5QEQPi6HiLGsO7UmkpfyaYNQ3D+rN6w6WfwH+nuqcGXWvawGwxOEroO4YGnFh95azw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/freebsd-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/freebsd-arm64/-/freebsd-arm64-0.28.2.tgz", + "integrity": "sha512-n+I0BTSRIoy+d6RPKnEVwql5UwBJolytvY4mAOIEJorKlqgPII8ix6slVVrfZ5Tnj7glIZvloylbB/EJPMWEXw==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/freebsd-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/freebsd-x64/-/freebsd-x64-0.28.2.tgz", + "integrity": "sha512-78XJTJkvPs0kz2w61301PJjXl4g7q3JqiYMZ/M/yVI73EHBrCRTgkhu9oqG7vPqq+a/yadEW8aD+agKlk5xrmg==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-arm": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-arm/-/linux-arm-0.28.2.tgz", + "integrity": "sha512-XlDnu2q5yoqems+xay6wSAcg9DDD7K9RLKZEBOMZm3ckNpJBvOX20tSfby8KfrrhINDyv9V2YVZKY/SpoGJI8w==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-arm64/-/linux-arm64-0.28.2.tgz", + "integrity": "sha512-pW4AC0P3it8c7do9MVM4p51FzHzdM/TZrerurgRcHJ2WTa1VQ1CIq18xncfpBJw4ojkiZZrKW2yIBWBP92j6Ug==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-ia32": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-ia32/-/linux-ia32-0.28.2.tgz", + "integrity": "sha512-CYbnj78HsIeA+DhgUKgFCfvNsTHFhMMrinUrMZpDXJXKN8T3XViTZ/+wtHeVxEWY8ewSzTFN+nRmSwO2tZaLUQ==", + "cpu": [ + "ia32" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-loong64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-loong64/-/linux-loong64-0.28.2.tgz", + "integrity": "sha512-buwkd8nsph4R+ajRvw0qM5Hja/TXQow3ptzWO2EbG/cqcIkHloRrdlBtQlshyYGTNFvfkfJ5tpPLVkY4DtsPfQ==", + "cpu": [ + "loong64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-mips64el": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-mips64el/-/linux-mips64el-0.28.2.tgz", + "integrity": "sha512-ZVykbDyk7519VwiNb9Lcj9m8XM6v5V9uKPvrEMkkEedVewf+0itkhahp4HDpgERXhwLRpWFypsGbG/J8s0QjJA==", + "cpu": [ + "mips64el" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-ppc64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-ppc64/-/linux-ppc64-0.28.2.tgz", + "integrity": "sha512-CAXl+Dtd9UUuJd8pKKdwh6MLm3MUMiqMPmhZ3tTSXPqfyQ3vDl6R5hZdZ/kYojK4ofXtdfSv1tFq8XzWx3heNQ==", + "cpu": [ + "ppc64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-riscv64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-riscv64/-/linux-riscv64-0.28.2.tgz", + "integrity": "sha512-GeXCej4IQtU1B+QlDV8W/RRvbzI3O/Stss+/bCXv4lZls5WGRtu2a+3JkA3i4qIUlMXpcHebWpF8AkJhATowuA==", + "cpu": [ + "riscv64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-s390x": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-s390x/-/linux-s390x-0.28.2.tgz", + "integrity": "sha512-3H1weTYZPxt/WOhByszQZybS9w5lKzUn1FDMsgEChbHWQwHYQQRfBxgCcZvPhjHfKyJjIievvMmEUawJrdY9Dg==", + "cpu": [ + "s390x" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/linux-x64/-/linux-x64-0.28.2.tgz", + "integrity": "sha512-4xTZr1FUmSoQW4XIWmit3tzQrUTZM+N3P0XV8xROKYF50XfI7xeO90+1bZvNwxIufQ9hDQVRJH5YhgPVF8A/HQ==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/netbsd-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/netbsd-arm64/-/netbsd-arm64-0.28.2.tgz", + "integrity": "sha512-sSATRjPeDBg3pdgHoQfoYBob11Kk1FGa9lui5RIHZCoCkJa9QKlvl3/vKz2usCmYYjs7ymJR/2Nnsqe+Hjt5nw==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "netbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/netbsd-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/netbsd-x64/-/netbsd-x64-0.28.2.tgz", + "integrity": "sha512-lqnzCV+mM0gIADaKihiCg6ifgfU2L3h5E33rNQBN1Y4MaVGnzryzmvvf7UHxprpQdE8hpqLolJ9Rl+SkIRDpyw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "netbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/openbsd-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/openbsd-arm64/-/openbsd-arm64-0.28.2.tgz", + "integrity": "sha512-AL2qJILH7lNjrDmCQDvdxMfAUIv8KMNZOvrwAQ8i8//ntL9FflhOyMJ8OZSMBb8/AWXe3/5v5S20y3zCoZWKoQ==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "openbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/openbsd-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/openbsd-x64/-/openbsd-x64-0.28.2.tgz", + "integrity": "sha512-QtiuPytchRyC4rwUKhexJdQKvDuZ6hWloi3igqPQNUJCS1/v9EiO3UTOXR6A3FoMo4fnAKbWJdqaIwhOzh8qEw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "openbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/openharmony-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/openharmony-arm64/-/openharmony-arm64-0.28.2.tgz", + "integrity": "sha512-WkhYDmpTjLvGlScA1rwjRUmhl4k8oXR3cIbtqWmELgU/dFeHHlEllxDvdWcNJV9rbzCexB5vz8gtNewWLgCT7Q==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "openharmony" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/sunos-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/sunos-x64/-/sunos-x64-0.28.2.tgz", + "integrity": "sha512-GPMSkTOtMnv2U2F8gxe4Io6qmVs+YKyp832Etqqxr0hFngmXQ3rzwytelm3GIn7T4VviRUlf3sOgBOiTdvaf7g==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "sunos" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/win32-arm64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/win32-arm64/-/win32-arm64-0.28.2.tgz", + "integrity": "sha512-PIhhEkE9uPBleRBrQEJpUn7MBnibZzbGzYWPmY3x+YoVg/95zbjB4CxPPOQ8l5tYYM4mMaCthF8/1DIfBQQyWQ==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/win32-ia32": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/win32-ia32/-/win32-ia32-0.28.2.tgz", + "integrity": "sha512-YmJbfTlvU7Sdn9BB+4PRES4oB6pxgS37MAONj+hBr/cpXS1aBPKXxNnDbu+QCWPj0o9dgyxeq79g6c5P8KeuYA==", + "cpu": [ + "ia32" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/win32-x64": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/@esbuild/win32-x64/-/win32-x64-0.28.2.tgz", + "integrity": "sha512-5ebpxr3nWMzrL/rnUI755Jkuee0bHL/Gq0WTF9lvcpv73wAp5eu8MfBUgWK9bhWvZjj7yX8etf/8tI8Ney695g==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@types/node": { + "version": "22.20.5", + "resolved": "https://registry.npmjs.org/@types/node/-/node-22.20.5.tgz", + "integrity": "sha512-U2+DNr+wSjpsTS/wZGYHq7GcwfuSmKiKvoPvK22zwTlRhU91yOniN4qRR5KhIjvif7ysw/dz/hKmfDH0Ris4aA==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, + "node_modules/@types/vscode": { + "version": "1.125.0", + "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.125.0.tgz", + "integrity": "sha512-0icm/ZQAaism87P0ekHqi4/Ju9du+Tm0RUW+y7vqRsxY2cY0FNRX1nAnaW7nT6npPt2tfHiheZ55Zm9UhqonFA==", + "dev": true, + "license": "MIT" + }, + "node_modules/esbuild": { + "version": "0.28.2", + "resolved": "https://registry.npmjs.org/esbuild/-/esbuild-0.28.2.tgz", + "integrity": "sha512-HKVLS8dvII+xoKW9kmqxbRKrnWEXfJJr/FZhhJmiqIB0e053QNYFqOBouTMO/k5sID4MvCiUCvv8b9M4h32wIA==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "bin": { + "esbuild": "bin/esbuild" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "@esbuild/aix-ppc64": "0.28.2", + "@esbuild/android-arm": "0.28.2", + "@esbuild/android-arm64": "0.28.2", + "@esbuild/android-x64": "0.28.2", + "@esbuild/darwin-arm64": "0.28.2", + "@esbuild/darwin-x64": "0.28.2", + "@esbuild/freebsd-arm64": "0.28.2", + "@esbuild/freebsd-x64": "0.28.2", + "@esbuild/linux-arm": "0.28.2", + "@esbuild/linux-arm64": "0.28.2", + "@esbuild/linux-ia32": "0.28.2", + "@esbuild/linux-loong64": "0.28.2", + "@esbuild/linux-mips64el": "0.28.2", + "@esbuild/linux-ppc64": "0.28.2", + "@esbuild/linux-riscv64": "0.28.2", + "@esbuild/linux-s390x": "0.28.2", + "@esbuild/linux-x64": "0.28.2", + "@esbuild/netbsd-arm64": "0.28.2", + "@esbuild/netbsd-x64": "0.28.2", + "@esbuild/openbsd-arm64": "0.28.2", + "@esbuild/openbsd-x64": "0.28.2", + "@esbuild/openharmony-arm64": "0.28.2", + "@esbuild/sunos-x64": "0.28.2", + "@esbuild/win32-arm64": "0.28.2", + "@esbuild/win32-ia32": "0.28.2", + "@esbuild/win32-x64": "0.28.2" + } + }, + "node_modules/typescript": { + "version": "5.9.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", + "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + } + } +} diff --git a/vscode-extension/package.json b/vscode-extension/package.json new file mode 100644 index 000000000..8bc404b46 --- /dev/null +++ b/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" + } +} diff --git a/vscode-extension/src/auth.ts b/vscode-extension/src/auth.ts new file mode 100644 index 000000000..d5ac02555 --- /dev/null +++ b/vscode-extension/src/auth.ts @@ -0,0 +1,73 @@ +import { exec } from "node:child_process"; + +// Fetch a new token this long before the current one expires. +const EXPIRY_MARGIN_MS = 5 * 60_000; +// Tokens without a readable `exp` (a PAT, say) are fetched again after an hour. +const DEFAULT_LIFETIME_MS = 60 * 60_000; + +const run = (command: string): Promise => + new Promise((resolve, reject) => { + exec(command, { timeout: 60_000, windowsHide: true }, (error, stdout, stderr) => { + if (error) { + reject(new Error(`Couldn't get a Databricks token: \`${command}\` failed: ${String(stderr).trim().slice(0, 500) || error.message}`)); + } else { + resolve(String(stdout).trim()); + } + }); + }); + +export const expiresAt = (token: string): number => { + try { + const { exp } = JSON.parse(Buffer.from(token.split(".")[1] ?? "", "base64url").toString()); + if (typeof exp === "number") { + return exp * 1000; + } + } catch { + // Not a JWT. + } + return Date.now() + DEFAULT_LIFETIME_MS; +}; + +/** Runs the auth command for a token and keeps it until it nears expiry or is rejected. */ +export class Tokens { + private cached: { value: string; expiresAt: number } | undefined; + private pending: Promise | undefined; + + constructor( + private readonly command: string, + private readonly onFetch: (expiresAt: number) => void = () => {}, + ) {} + + get(): Promise { + if (this.cached && this.cached.expiresAt - EXPIRY_MARGIN_MS > Date.now()) { + return Promise.resolve(this.cached.value); + } + return this.fetch(this.command); + } + + /** A token the gateway rejected. The CLI can hand back its cached copy while `exp` is still ahead, so force a new one. */ + async renew(rejected: string): Promise { + await this.pending?.catch(() => {}); + if (this.cached && this.cached.value !== rejected) { + return this.cached.value; + } + this.cached = undefined; + return this.fetch(`${this.command} --force-refresh`); + } + + private fetch(command: string): Promise { + this.pending ??= run(command) + .then((value) => { + if (!value) { + throw new Error(`Couldn't get a Databricks token: \`${command}\` printed nothing.`); + } + this.cached = { value, expiresAt: expiresAt(value) }; + this.onFetch(this.cached.expiresAt); + return value; + }) + .finally(() => { + this.pending = undefined; + }); + return this.pending; + } +} diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts new file mode 100644 index 000000000..991e27485 --- /dev/null +++ b/vscode-extension/src/extension.ts @@ -0,0 +1,139 @@ +import * as vscode from "vscode"; +import { version } from "../package.json"; +import { Tokens } from "./auth"; +import { chatCompletionsReplies, chatCompletionsRequest, estimateTokens, responsesReplies, responsesRequest, type Reply } from "./formats"; +import { streamEvents } from "./gateway"; + +export interface ModelSetting { + readonly id: string; + readonly api: "chat_completions" | "responses"; + readonly maxInputTokens?: number; + readonly maxOutputTokens?: number; +} + +export interface Settings { + readonly workspace: string; + readonly models: readonly ModelSetting[]; + readonly authCommand: string; +} + +const tokenLimits = (model: ModelSetting): Partial => + Object.fromEntries( + (["maxInputTokens", "maxOutputTokens"] as const).flatMap((key) => { + const value: unknown = model[key]; + return typeof value === "number" && Number.isInteger(value) && value > 0 ? [[key, value]] : []; + }), + ); + +export const readSettings = (get: (key: string) => unknown): Settings => { + const models = get("models"); + return { + workspace: String(get("workspace") ?? "").replace(/\/+$/, ""), + models: (Array.isArray(models) ? models : []) + .filter( + (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) })), + authCommand: String(get("authCommand") ?? ""), + }; +}; + +/** The gateway attributes requests by this User-Agent. */ +export const userAgent = (copilotChatVersion: unknown): string => + `ucode/${version} githubcopilotchat${typeof copilotChatVersion === "string" ? `/${copilotChatVersion}` : ""}`; + +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( + private readonly settings: () => Settings, + private readonly log: (line: string) => void, + 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 }) => ({ + id, + api, + 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, + detail: "Databricks", + // The gateway reports no limits, so a model ug configure has none for gets Copilot's 128K default. + maxInputTokens: maxInputTokens ?? 128_000, + maxOutputTokens: maxOutputTokens ?? 32_000, + capabilities: { toolCalling: true, imageInput: api === "responses" || id.includes("claude") }, + })); + } + + async provideLanguageModelChatResponse( + model: Model, + messages: readonly vscode.LanguageModelChatRequestMessage[], + options: vscode.ProvideLanguageModelChatResponseOptions, + progress: vscode.Progress, + token: vscode.CancellationToken, + ): Promise { + const { workspace, authCommand } = this.settings(); + this.tokens ??= new Tokens(authCommand, (expiresAt) => this.log(`Fetched a Databricks token (expires ${new Date(expiresAt).toISOString()})`)); + const turn = { model: model.id, messages, tools: options.tools ?? [], toolRequired: options.toolMode === vscode.LanguageModelChatToolMode.Required }; + const abort = new AbortController(); + const cancelled = token.onCancellationRequested(() => abort.abort()); + try { + const replies: AsyncIterable = + model.api === "responses" + ? responsesReplies(streamEvents(workspace, "responses", responsesRequest(turn), this.tokens, this.headers(), abort.signal)) + : chatCompletionsReplies(streamEvents(workspace, "chat/completions", chatCompletionsRequest(turn), this.tokens, this.headers(), abort.signal)); + for await (const reply of replies) { + progress.report(reply); + } + } catch (error) { + if (!token.isCancellationRequested) { + this.log(`${model.id} failed: ${error instanceof Error ? error.message : String(error)}`); + throw error; + } + } finally { + cancelled.dispose(); + } + } + + async provideTokenCount(_model: Model, input: string | vscode.LanguageModelChatRequestMessage): Promise { + return estimateTokens(input); + } +} + +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)); + }; + 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(); + } + }), + ); +} diff --git a/vscode-extension/src/formats.ts b/vscode-extension/src/formats.ts new file mode 100644 index 000000000..d2c61f7e9 --- /dev/null +++ b/vscode-extension/src/formats.ts @@ -0,0 +1,207 @@ +import * as vscode from "vscode"; + +// The stable API names only User and Assistant; Copilot Chat sends its system prompt with role 3. +const SYSTEM_ROLE = 3; + +export interface Turn { + readonly model: string; + readonly messages: readonly vscode.LanguageModelChatRequestMessage[]; + readonly tools: readonly vscode.LanguageModelChatTool[]; + readonly toolRequired: boolean; +} + +export type Reply = vscode.LanguageModelTextPart | vscode.LanguageModelToolCallPart; + +const IMAGE_TOKENS = 1000; + +const textOf = (part: unknown): string => { + if (part instanceof vscode.LanguageModelTextPart) { + return part.value; + } + if (part instanceof vscode.LanguageModelDataPart && !part.mimeType.startsWith("image/")) { + return new TextDecoder().decode(part.data); + } + if (part instanceof vscode.LanguageModelPromptTsxPart) { + return JSON.stringify(part.value); + } + return ""; +}; + +const imageUrl = (part: unknown): string | undefined => + part instanceof vscode.LanguageModelDataPart && part.mimeType.startsWith("image/") + ? `data:${part.mimeType};base64,${Buffer.from(part.data).toString("base64")}` + : undefined; + +// The gateway requires a schema on every function tool; VS Code omits it for tools that take no input. +const parametersOf = (tool: vscode.LanguageModelChatTool): object => tool.inputSchema ?? { type: "object", properties: {} }; + +const toolCall = (callId: string, name: string, args: string): vscode.LanguageModelToolCallPart => { + let input: unknown; + try { + input = args.trim() ? JSON.parse(args) : {}; + } catch { + input = undefined; + } + if (typeof input !== "object" || input === null || Array.isArray(input)) { + throw new Error(`The model called ${name} with arguments that aren't a JSON object: ${args}`); + } + return new vscode.LanguageModelToolCallPart(callId, name, input); +}; + +const outputLimit = () => new Error("The model stopped at its output token limit before finishing the response."); + +const toolOptions = (turn: Turn, toTool: (tool: vscode.LanguageModelChatTool) => object): object => + turn.tools.length ? { tools: turn.tools.map(toTool), ...(turn.toolRequired ? { tool_choice: "required" } : {}) } : {}; + +/** A rough count: about four characters per token, and a flat cost per image. */ +export const estimateTokens = (input: string | vscode.LanguageModelChatRequestMessage): number => { + if (typeof input === "string") { + return Math.ceil(input.length / 4); + } + let characters = 0; + let images = 0; + for (const part of input.content) { + const items = part instanceof vscode.LanguageModelToolResultPart ? part.content : [part]; + for (const item of items) { + characters += item instanceof vscode.LanguageModelToolCallPart ? item.name.length + JSON.stringify(item.input).length : textOf(item).length; + images += imageUrl(item) === undefined ? 0 : 1; + } + } + return Math.ceil(characters / 4) + images * IMAGE_TOKENS; +}; + +/** 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), +}); + +export function chatCompletionsRequest(turn: Turn): object { + const messages = turn.messages.flatMap((message): object[] => { + const { results, calls, text, images } = sortParts(message.content); + if (Number(message.role) === SYSTEM_ROLE) { + return [{ role: "system", content: text }]; + } + if (message.role === vscode.LanguageModelChatMessageRole.Assistant) { + const toolCalls = calls.map((call) => ({ + id: call.callId, + type: "function", + function: { name: call.name, arguments: JSON.stringify(call.input) }, + })); + return text || toolCalls.length ? [{ role: "assistant", content: text || null, ...(toolCalls.length ? { tool_calls: toolCalls } : {}) }] : []; + } + const content = [ + ...(text ? [{ type: "text", text }] : []), + ...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("") })), + ...(content.length ? [{ role: "user", content }] : []), + ]; + }); + return { + model: turn.model, + messages, + stream: true, + ...toolOptions(turn, (tool) => ({ type: "function", function: { name: tool.name, description: tool.description, parameters: parametersOf(tool) } })), + }; +} + +type ChatContent = string | { type?: string; text?: string }[] | null; + +interface ChatChunk { + choices?: { + delta?: { content?: ChatContent; tool_calls?: { index?: number; id?: string; function?: { name?: string; arguments?: string } }[] }; + finish_reason?: string | null; + }[]; +} + +/** While a Claude model thinks, the gateway sends `content` as blocks rather than a string. Only the + * text ones are the reply: passing a block along would make Copilot Chat fail the turn. */ +const deltaText = (content: ChatContent | undefined): string => + typeof content === "string" + ? content + : (content ?? []) + .filter((block) => block.type === "text") + .map((block) => block.text ?? "") + .join(""); + +export async function* chatCompletionsReplies(events: AsyncIterable): AsyncGenerator { + const calls: { id: string; name: string; args: string }[] = []; + for await (const event of events) { + for (const choice of (event as ChatChunk).choices ?? []) { + const text = deltaText(choice.delta?.content); + if (text) { + yield new vscode.LanguageModelTextPart(text); + } + for (const fragment of choice.delta?.tool_calls ?? []) { + const call = (calls[fragment.index ?? 0] ??= { id: "", name: "", args: "" }); + call.id = fragment.id || call.id; + call.name = fragment.function?.name || call.name; + call.args += fragment.function?.arguments ?? ""; + } + if (choice.finish_reason === "length") { + throw outputLimit(); + } + } + } + for (const call of calls.filter(Boolean)) { + yield toolCall(call.id, call.name, call.args); + } +} + +// 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[] => { + const { results, calls, text, images } = sortParts(message.content); + if (Number(message.role) === SYSTEM_ROLE) { + return [{ role: "system", content: text }]; + } + if (message.role === vscode.LanguageModelChatMessageRole.Assistant) { + return [ + ...(text ? [{ role: "assistant", content: text }] : []), + ...calls.map((call) => ({ type: "function_call", call_id: call.callId, name: call.name, arguments: JSON.stringify(call.input) })), + ]; + } + const content = [ + ...(text ? [{ type: "input_text", text }] : []), + ...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("") })), + ...(content.length ? [{ role: "user", content }] : []), + ]; + }); + return { + model: turn.model, + input, + stream: true, + store: false, + ...toolOptions(turn, (tool) => ({ type: "function", name: tool.name, description: tool.description, parameters: parametersOf(tool), strict: false })), + }; +} + +interface ResponseEvent { + type: string; + delta?: string; + item?: { type?: string; call_id?: string; name?: string; arguments?: string }; + response?: { incomplete_details?: { reason?: string } | null; error?: { message?: string } | null }; + message?: string; +} + +export async function* responsesReplies(events: AsyncIterable): AsyncGenerator { + 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(); + } else if (event.type === "response.failed" || event.type === "error") { + throw new Error(event.response?.error?.message ?? event.message ?? "The model response failed."); + } + } +} diff --git a/vscode-extension/src/gateway.ts b/vscode-extension/src/gateway.ts new file mode 100644 index 000000000..96d5eafdd --- /dev/null +++ b/vscode-extension/src/gateway.ts @@ -0,0 +1,55 @@ +import type { Tokens } from "./auth"; + +/** POSTs a streaming request to the workspace's AI Gateway and yields each server-sent event's JSON. */ +export async function* streamEvents( + workspace: string, + path: "chat/completions" | "responses", + body: object, + tokens: Tokens, + headers: Record, + signal: AbortSignal, +): AsyncGenerator { + const post = (token: string) => + fetch(`${workspace}/ai-gateway/mlflow/v1/${path}`, { + method: "POST", + headers: { ...headers, Authorization: `Bearer ${token}`, "Content-Type": "application/json" }, + body: JSON.stringify(body), + signal, + }); + let token = await tokens.get(); + let response = await post(token); + if (response.status === 401) { + await response.body?.cancel(); + token = await tokens.renew(token); + response = await post(token); + } + if (response.status === 401) { + await response.body?.cancel(); + throw new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${workspace}\`, then retry.`); + } + if (!response.ok || !response.body) { + throw new Error(`AI Gateway returned ${response.status}: ${(await response.text()).slice(0, 300)}`); + } + const decoder = new TextDecoder(); + let buffer = ""; + for await (const chunk of response.body) { + buffer += decoder.decode(chunk, { stream: true }); + const lines = buffer.split("\n"); + buffer = lines.pop() ?? ""; + for (const line of lines) { + const data = dataOf(line); + if (data === "[DONE]") { + return; + } + if (data) { + yield JSON.parse(data); + } + } + } + const last = dataOf(buffer + decoder.decode()); + if (last && last !== "[DONE]") { + yield JSON.parse(last); + } +} + +const dataOf = (line: string): string => (line.startsWith("data:") ? line.slice(5).trim() : ""); diff --git a/vscode-extension/test/auth.test.ts b/vscode-extension/test/auth.test.ts new file mode 100644 index 000000000..484085c48 --- /dev/null +++ b/vscode-extension/test/auth.test.ts @@ -0,0 +1,51 @@ +import assert from "node:assert/strict"; +import { describe, it } from "node:test"; +import { expiresAt, Tokens } from "../src/auth"; +import { jwt, rejectsWith, tokenCommand } from "./support"; + +describe("Tokens", () => { + it("reuses a token until it nears expiry", async () => { + const valid = tokenCommand([jwt(3600)]); + const tokens = new Tokens(valid.command); + assert.equal(await tokens.get(), await tokens.get()); + assert.equal(valid.runs(), 1); + + const expiring = tokenCommand([jwt(60), jwt(3600)]); + const refreshed = new Tokens(expiring.command); + assert.notEqual(await refreshed.get(), await refreshed.get()); + assert.equal(expiring.runs(), 2); + }); + + it("runs the command once for concurrent requests", async () => { + const source = tokenCommand([jwt(3600)]); + const tokens = new Tokens(source.command); + await Promise.all([tokens.get(), tokens.get(), tokens.get()]); + assert.equal(source.runs(), 1); + }); + + it("forces a new token after a rejection, but not after a stale one", async () => { + const source = tokenCommand(["fresh"], "cached"); + const tokens = new Tokens(source.command); + assert.equal(await tokens.get(), "cached"); + assert.equal(await tokens.renew("cached"), "fresh"); + assert.equal(await tokens.renew("cached"), "fresh"); + assert.equal(source.runs(), 2); + }); + + it("waits for a fetch in flight before deciding to force one", async () => { + const source = tokenCommand(["fresh"], "cached"); + const tokens = new Tokens(source.command); + assert.deepEqual(await Promise.all([tokens.get(), tokens.renew("cached")]), ["cached", "fresh"]); + assert.equal(source.runs(), 2); + }); + + it("reports the command's failure and empty output", async () => { + await rejectsWith(new Tokens(`node -e "console.error('sign-in expired'); process.exit(1)"`).get(), "sign-in expired"); + await rejectsWith(new Tokens(`node -e ""`).get(), "printed nothing"); + }); + + it("reads exp from a JWT and gives other tokens an hour", () => { + assert.equal(expiresAt(`h.${Buffer.from('{"exp":1800000000}').toString("base64url")}.s`), 1_800_000_000_000); + assert.ok(expiresAt("dapi123") - Date.now() > 59 * 60_000); + }); +}); diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts new file mode 100644 index 000000000..5e94cfcf4 --- /dev/null +++ b/vscode-extension/test/extension.test.ts @@ -0,0 +1,167 @@ +import assert from "node:assert/strict"; +import { existsSync } from "node:fs"; +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 { streamEvents } from "../src/gateway"; +import { collect, message, rejectsWith, text, tokenCommand, USER } from "./support"; +import { LanguageModelChatToolMode, LanguageModelTextPart } from "./vscode"; + +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 }[] = []; + 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)) }); + 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); + const middle = Math.floor(bytes.length / 2); + return new Response(new Blob([bytes.slice(0, middle), bytes.slice(middle)]).stream(), { status }); + }) as typeof fetch; + return requests; +}; + +const realFetch = globalThis.fetch; +afterEach(() => { + globalThis.fetch = realFetch; +}); + +describe("streamEvents", () => { + it("yields each event's JSON until [DONE]", async () => { + stubFetch([200, sse('{"a":1}', '{"b":2}', "[DONE]", '{"c":3}')]); + const tokens = new Tokens(tokenCommand(["t"]).command); + assert.deepEqual(await collect(streamEvents("https://ws", "responses", {}, tokens, {}, new AbortController().signal)), [{ a: 1 }, { b: 2 }]); + }); + + it("keeps a last event that has no trailing newline", async () => { + stubFetch([200, `data: {"a":1}\n\ndata: {"b":2}`]); + const tokens = new Tokens(tokenCommand(["t"]).command); + assert.deepEqual(await collect(streamEvents("https://ws", "responses", {}, tokens, {}, new AbortController().signal)), [{ a: 1 }, { b: 2 }]); + }); + + it("retries once with a new token after a 401, then asks the user to sign in", async () => { + let requests = stubFetch([401, ""], [200, sse("[DONE]")]); + await collect(streamEvents("https://ws", "chat/completions", {}, new Tokens(tokenCommand(["fresh"], "rejected").command), {}, new AbortController().signal)); + assert.deepEqual(requests.map((request) => request.authorization), ["Bearer rejected", "Bearer fresh"]); + assert.equal(requests[0]?.url, "https://ws/ai-gateway/mlflow/v1/chat/completions"); + + requests = stubFetch([401, ""]); + await rejectsWith( + collect(streamEvents("https://ws", "responses", {}, new Tokens(tokenCommand(["a", "b"]).command), {}, new AbortController().signal)), + "Run `databricks auth login --host https://ws`", + ); + assert.equal(requests.length, 2); + }); + + it("reports other gateway errors with their body", async () => { + stubFetch([400, '{"message":"bad tools"}']); + await rejectsWith( + collect(streamEvents("https://ws", "responses", {}, new Tokens(tokenCommand(["t"]).command), {}, new AbortController().signal)), + 'AI Gateway returned 400: {"message":"bad tools"}', + ); + }); +}); + +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", () => { + it("ships the icon it names", () => { + assert.ok(packageJson.contributes.languageModelChatProviders.every((provider) => existsSync(provider.icon))); + }); +}); + +describe("userAgent", () => { + it("names ug's extension and Copilot Chat's version", () => { + assert.equal(userAgent("0.71.0"), `ucode/${packageJson.version} githubcopilotchat/0.71.0`); + assert.equal(userAgent(undefined), `ucode/${packageJson.version} githubcopilotchat`); + }); +}); + +describe("readSettings", () => { + it("keeps well-formed models and token limits and trims the workspace URL", () => { + const values: Record = { + workspace: "https://ws/", + authCommand: "ug auth-token", + models: [ + { id: "a", api: "responses" }, + { 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 }, + ], + }; + assert.deepEqual(readSettings((key) => values[key]), { + workspace: "https://ws", + authCommand: "ug auth-token", + models: [ + { id: "a", api: "responses" }, + { id: "c", api: "chat_completions", maxInputTokens: 936_000 }, + { id: "d", api: "chat_completions" }, + ], + }); + }); +}); + +describe("DatabricksModels", () => { + const settings = (command: string): Settings => ({ + 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.gpt-6-luna", api: "responses" }, + ], + }); + + it("lists the configured models, 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]), + [ + ["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], + ], + ); + assert.deepEqual(await new DatabricksModels(() => ({ ...settings(""), models: [] }), (line) => logs.push(line)).provideLanguageModelChatInformation(), []); + assert.deepEqual(logs, ["No Databricks models: run `ug configure --agents vscode` to set them up."]); + }); + + it("sends each model's turns to its API with the User-Agent", async () => { + const requests = stubFetch( + [200, sse('{"choices":[{"delta":{"content":"hello"}}]}', "[DONE]")], + [200, sse('{"type":"response.output_text.delta","delta":"hi"}')], + ); + const provider = new DatabricksModels(() => settings(tokenCommand(["t"]).command), () => {}, () => ({ "User-Agent": "ucode/0 githubcopilotchat" })); + const [claude, gpt] = await provider.provideLanguageModelChatInformation(); + const reported: unknown[] = []; + for (const model of [claude!, gpt!]) { + await provider.provideLanguageModelChatResponse( + model, + [message(USER, text("hi"))], + { toolMode: LanguageModelChatToolMode.Auto } as unknown as vscode.ProvideLanguageModelChatResponseOptions, + { report: (part) => reported.push(part) }, + { isCancellationRequested: false, onCancellationRequested: () => ({ dispose() {} }) } as unknown as vscode.CancellationToken, + ); + } + assert.deepEqual( + requests.map((request) => [request.url, request.body.model]), + [ + ["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"], + ], + ); + 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 new file mode 100644 index 000000000..9f66d8817 --- /dev/null +++ b/vscode-extension/test/formats.test.ts @@ -0,0 +1,122 @@ +import assert from "node:assert/strict"; +import { describe, it } from "node:test"; +import { chatCompletionsReplies, chatCompletionsRequest, estimateTokens, responsesReplies, responsesRequest } from "../src/formats"; +import { ASSISTANT, call, collect, events, image, message, rejectsWith, result, SYSTEM, text, USER, weatherTool } from "./support"; +import { LanguageModelDataPart, LanguageModelTextPart, LanguageModelToolCallPart } from "./vscode"; + +const conversation = [ + message(SYSTEM, text("Be brief.")), + message(USER, text("Weather?"), image(7)), + message(ASSISTANT, text("Checking."), call("c1", "get_weather", { city: "Paris" })), + message(USER, result("c1", text("Sunny"))), +]; +const turn = { model: "m", messages: conversation, tools: [weatherTool], toolRequired: true }; + +describe("Chat Completions", () => { + it("sends the conversation, images and tools", () => { + assert.deepEqual(chatCompletionsRequest(turn), { + model: "m", + stream: true, + messages: [ + { role: "system", content: "Be brief." }, + { role: "user", content: [{ type: "text", text: "Weather?" }, { type: "image_url", image_url: { url: "data:image/png;base64,Bw==" } }] }, + { role: "assistant", content: "Checking.", tool_calls: [{ id: "c1", type: "function", function: { name: "get_weather", arguments: '{"city":"Paris"}' } }] }, + { role: "tool", tool_call_id: "c1", content: "Sunny" }, + ], + tools: [{ type: "function", function: { name: "get_weather", description: "Weather for a city", parameters: weatherTool.inputSchema } }], + tool_choice: "required", + }); + }); + + 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, [ + { type: "function", function: { name: "list", description: "List", parameters: { type: "object", properties: {} } } }, + ]); + assert.equal("tool_choice" in noSchema, false); + assert.equal("tools" in chatCompletionsRequest({ ...turn, tools: [] }), false); + }); + + it("streams text and assembles a tool call from its fragments", async () => { + const replies = await collect( + chatCompletionsReplies( + events( + { choices: [{ delta: { content: "Checking." } }] }, + { choices: [{ delta: { tool_calls: [{ index: 0, id: "c1", function: { name: "get_weather", arguments: '{"ci' } }] } }] }, + { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: 'ty":"Paris"}' } }] }, finish_reason: "tool_calls" }] }, + ), + ), + ); + assert.deepEqual(replies, [new LanguageModelTextPart("Checking."), new LanguageModelToolCallPart("c1", "get_weather", { city: "Paris" })]); + }); + + it("takes the text out of content blocks and drops a thinking model's reasoning", async () => { + const replies = await collect( + 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." }] } }] }, + ), + ), + ); + assert.deepEqual(replies, [new LanguageModelTextPart("Checking. Done.")]); + }); + + 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]" } }] } }] }))), + "aren't a JSON object", + ); + }); +}); + +describe("estimateTokens", () => { + it("counts text and gives each image a flat cost instead of its bytes", () => { + const screenshot = new Uint8Array(400_000); + const withImage = message(USER, text("abcd".repeat(10)), result("c1", text("ok"), image(1)), new LanguageModelDataPart(screenshot, "image/png")); + assert.equal(estimateTokens("abcdefgh"), 2); + assert.equal(estimateTokens(withImage), 11 + 2000); + }); +}); + +describe("Responses", () => { + it("sends the whole conversation statelessly", () => { + assert.deepEqual(responsesRequest(turn), { + model: "m", + stream: true, + store: false, + input: [ + { role: "system", content: "Be brief." }, + { role: "user", content: [{ type: "input_text", text: "Weather?" }, { type: "input_image", image_url: "data:image/png;base64,Bw==", detail: "auto" }] }, + { role: "assistant", content: "Checking." }, + { type: "function_call", call_id: "c1", name: "get_weather", arguments: '{"city":"Paris"}' }, + { type: "function_call_output", call_id: "c1", output: "Sunny" }, + ], + tools: [{ type: "function", name: "get_weather", description: "Weather for a city", parameters: weatherTool.inputSchema, strict: false }], + tool_choice: "required", + }); + }); + + it("streams text and each finished function call", async () => { + const replies = await collect( + responsesReplies( + events( + { type: "response.output_text.delta", delta: "Checking." }, + { type: "response.function_call_arguments.delta", delta: '{"city"' }, + { type: "response.output_item.done", item: { type: "function_call", call_id: "c1", name: "get_weather", arguments: '{"city":"Paris"}' } }, + { type: "response.completed", response: {} }, + ), + ), + ); + assert.deepEqual(replies, [new LanguageModelTextPart("Checking."), new LanguageModelToolCallPart("c1", "get_weather", { city: "Paris" })]); + }); + + it("fails on a token-limit stop and on a failed response", async () => { + await rejectsWith( + collect(responsesReplies(events({ type: "response.incomplete", response: { incomplete_details: { reason: "max_output_tokens" } } }))), + "output token limit", + ); + await rejectsWith(collect(responsesReplies(events({ type: "response.failed", response: { error: { message: "overloaded" } } }))), "overloaded"); + }); +}); diff --git a/vscode-extension/test/support.ts b/vscode-extension/test/support.ts new file mode 100644 index 000000000..7519d648e --- /dev/null +++ b/vscode-extension/test/support.ts @@ -0,0 +1,62 @@ +import assert from "node:assert/strict"; +import { mkdtempSync, readFileSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type * as vscode from "vscode"; +import { LanguageModelDataPart, LanguageModelTextPart, LanguageModelToolCallPart, LanguageModelToolResultPart } from "./vscode"; + +export const SYSTEM = 3; +export const USER = 1; +export const ASSISTANT = 2; + +export const message = (role: number, ...content: unknown[]) => ({ role, content, name: undefined }) as unknown as vscode.LanguageModelChatRequestMessage; +export const text = (value: string) => new LanguageModelTextPart(value); +export const image = (byte: number) => new LanguageModelDataPart(new Uint8Array([byte]), "image/png"); +export const call = (id: string, name: string, input: object) => new LanguageModelToolCallPart(id, name, input); +export const result = (id: string, ...content: unknown[]) => new LanguageModelToolResultPart(id, content); +export const weatherTool: vscode.LanguageModelChatTool = { + name: "get_weather", + description: "Weather for a city", + inputSchema: { type: "object", properties: { city: { type: "string" } } }, +}; + +export const jwt = (expiresInSeconds: number) => + `h.${Buffer.from(JSON.stringify({ exp: Math.floor(Date.now() / 1000) + expiresInSeconds })).toString("base64url")}.s`; + +/** + * A real shell command that prints the next of `tokens` on each run and counts its runs. With `cached`, it + * prints `cached` until it's run with --force-refresh, like the Databricks CLI's token cache. + */ +export const tokenCommand = (tokens: string[], cached?: string) => { + const dir = mkdtempSync(join(tmpdir(), "ug-vscode-test-")); + const state = join(dir, "runs"); + writeFileSync(state, "0"); + const script = join(dir, "token.js"); + writeFileSync( + script, + `const fs = require("fs"); const n = Number(fs.readFileSync(${JSON.stringify(state)}, "utf8")); fs.writeFileSync(${JSON.stringify(state)}, String(n + 1)); + const tokens = ${JSON.stringify(tokens)}; const cached = ${JSON.stringify(cached ?? null)}; + const token = cached !== null && !process.argv.includes("--force-refresh") ? cached : tokens[Math.min(n, tokens.length - 1)]; + process.stdout.write(token + "\\n");`, + ); + return { command: `node ${JSON.stringify(script)}`, runs: () => Number(readFileSync(state, "utf8")) }; +}; + +export async function* events(...items: unknown[]): AsyncGenerator { + yield* items; +} + +export const collect = async (items: AsyncIterable): Promise => { + const out: T[] = []; + for await (const item of items) { + out.push(item); + } + return out; +}; + +/** Asserts that `promise` rejects with an error whose message contains `text`. */ +export const rejectsWith = (promise: Promise, text: string) => + assert.rejects(promise, (error: Error) => { + assert.ok(error.message.includes(text), `expected "${error.message}" to contain "${text}"`); + return true; + }); diff --git a/vscode-extension/test/vscode.ts b/vscode-extension/test/vscode.ts new file mode 100644 index 000000000..ffb58110b --- /dev/null +++ b/vscode-extension/test/vscode.ts @@ -0,0 +1,40 @@ +// Unit tests run outside VS Code; this provides the parts of the `vscode` API the extension touches. +export class LanguageModelTextPart { + constructor(public value: string) {} +} +export class LanguageModelToolCallPart { + constructor( + public callId: string, + public name: string, + public input: object, + ) {} +} +export class LanguageModelToolResultPart { + constructor( + public callId: string, + public content: unknown[], + ) {} +} +export class LanguageModelDataPart { + constructor( + public data: Uint8Array, + public mimeType: string, + ) {} +} +export class LanguageModelPromptTsxPart { + constructor(public value: unknown) {} +} +export const LanguageModelChatMessageRole = { User: 1, Assistant: 2 } as const; +export const LanguageModelChatToolMode = { Auto: 1, Required: 2 } as const; +export class EventEmitter { + private listeners: ((value: T) => void)[] = []; + event = (listener: (value: T) => void) => { + this.listeners.push(listener); + return { dispose: () => {} }; + }; + fire(value: T): void { + for (const listener of this.listeners) { + listener(value); + } + } +} diff --git a/vscode-extension/tsconfig.json b/vscode-extension/tsconfig.json new file mode 100644 index 000000000..a12269ab6 --- /dev/null +++ b/vscode-extension/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "target": "ES2023", + "lib": ["ES2023", "DOM", "DOM.AsyncIterable"], + "module": "Preserve", + "moduleResolution": "Bundler", + "types": ["node"], + "strict": true, + "noUncheckedIndexedAccess": true, + "resolveJsonModule": true, + "isolatedModules": true, + "skipLibCheck": true, + "noEmit": true + }, + "include": ["src/**/*.ts", "test/**/*.ts"] +}