src/vs/platform/agentHost/node/copilot/byokLmProxyService.ts

285 LOC · 258 covered · 27 uncovered · 52 ranges · 932 concepts · 22 introducers · 435 tests

File neighbourhood

The centred file is linked to every concept that introduces one of its ranges, every test that runs code from the file, and the gray connector concepts standing between those tests and the file's own introducer concepts. Undirected links join concepts to every file where they introduce source and concepts to the tests they introduce; arrows show specialization between the displayed concepts and bridge only concepts omitted from this view. Concept colors match the source ranges below; connector concepts have no source color and are shown in gray.

Focused file, its introducer and connector concepts, their introduced files, and tests that run code from the file

In the embedded map, ordinary wheel input scrolls the page; use the visible controls to zoom and drag to pan. Open the full-screen map for canvas navigation: wheel pans, Ctrl/Command plus wheel zooms, and arrow keys pan when this region is focused. On touch screens, open the full-screen map to pan or pinch. If JavaScript or WebGL is unavailable, use the related-file, concept, and source links on this page.

Graph controls are ready.

Interactive rendering requires JavaScript and WebGL. Use the related-file, concept, and source links on this page while the interactive map is unavailable.

1 > /*--------------------------------------------------------------------------------------------- byokLmProxyService.ts ×11
2 > * Copyright (c) Microsoft Corporation. All rights reserved.
3 > * Licensed under the MIT License. See License.txt in the project root for license information.
4 > *--------------------------------------------------------------------------------------------*/
5 >
6 > import type * as http from 'http';
7 > import { createDecorator } from '../../../instantiation/common/instantiation.js';
8 > import { ILogService } from '../../../log/common/log.js';
9 > import { IByokLmBridgeRegistry } from '../byokLmBridgeRegistry.js';
10 > import { parseProxyBearer } from '../claude/claudeProxyAuth.js';
11 > import {
12 > ILoopbackProxyHandle,
13 > ILoopbackProxyRuntime,
14 > IProxyInFlight,
15 > LoopbackProxyServer,
16 > readProxyRequestBody,
17 > } from '../shared/loopbackProxyServer.js';
18 > import {
19 > IOpenAiChatRequest,
20 > OpenAiTranslationError,
21 > bridgeResultToSseFrames,
22 > openAiErrorBody,
23 > openAiRequestToBridge,
24 > } from './byokOpenAiTranslation.js';
25 >
26 > // #region Public types
27 >
28 > /**
29 > * Handle returned by {@link IByokLmProxyService.start}. Refcounts the shared
30 > * loopback server (see {@link LoopbackProxyServer}): when every handle is
31 > * disposed the listener closes and the nonce is destroyed; the next `start()`
32 > * rebinds with a fresh port and nonce.
33 > *
34 > * **Subprocess ownership invariant.** Callers that hand `baseUrl`/`nonce` to
35 > * the Copilot SDK runtime subprocess MUST kill that subprocess before calling
36 > * `dispose()` — after disposal the proxy may rebind on a different port and the
37 > * subprocess would silently lose its endpoint (same contract as the Claude and
38 > * Codex proxies).
39 > */
40 > export interface IByokLmProxyHandle extends ILoopbackProxyHandle {
41 > /** e.g. `http://127.0.0.1:54321` — no trailing slash. */
42 > readonly baseUrl: string;
43 > /** 256-bit hex string. Combine with a session id as `Bearer <nonce>.<sessionId>`. */
44 > readonly nonce: string;
45 > /**
46 > * Build the provider `baseUrl` for a given BYOK vendor. The vendor is
47 > * encoded into the path so a single proxy can serve every vendor; the
48 > * runtime appends `/chat/completions` to this URL.
49 > */
50 > providerBaseUrl(vendor: string): string;
51 > }
52 >
53 > export const IByokLmProxyService = createDecorator<IByokLmProxyService>('byokLmProxyService');
54 >
55 > export interface IByokLmProxyService {
56 > readonly _serviceBrand: undefined;
57 >
58 > /** Start the proxy (if not already running) and return a refcounted handle. */
59 > start(): Promise<IByokLmProxyHandle>;
60 >
61 > /**
62 > * Force-close the proxy regardless of refcount and abort in-flight
63 > * requests. Idempotent; subsequent `start()` calls rebind.
64 > */
65 > dispose(): void;
66 > }
67 >
68 > // #endregion
69 >
70 > const PROXY_USER_FACING_NAME = 'ByokLmProxyService';
71 > const VENDOR_PATH_PREFIX = '/v/';
72 > const CHAT_COMPLETIONS_SUFFIX = '/chat/completions';
73 >
74 > /**
75 > * The BYOK proxy keeps no per-bind mutable state: the active renderer bridge is
76 > * resolved from {@link IByokLmBridgeRegistry} at request time, and the nonce
77 > * lives on the runtime owned by {@link LoopbackProxyServer}.
78 > */
79 > type ByokLmProxyState = undefined;
80 >
81 > /**
82 > * Local OpenAI-compatible HTTP proxy that lets the Copilot SDK runtime run
83 > * BYOK models provided by VS Code extensions. The runtime is configured with a
84 > * `type: 'openai'`, `wireApi: 'completions'` provider whose `baseUrl` points
85 > * here; inbound `POST /v/<vendor>/chat/completions` requests are authenticated,
86 > * translated, and forwarded to the renderer LM API via
87 > * {@link IByokLmBridgeRegistry}, and the buffered completion is streamed back
88 > * as OpenAI Chat Completions SSE.
89 > *
90 > * The server lifecycle — lazy bind on `127.0.0.1`, nonce minting, refcounted
91 > * handles, in-flight tracking, and teardown — is inherited from
92 > * {@link LoopbackProxyServer}; this subclass only implements request routing.
93 > */
94 > export class ByokLmProxyService extends LoopbackProxyServer<ByokLmProxyState> implements IByokLmProxyService {
95 >
96 > declare readonly _serviceBrand: undefined;
97 >
98 > constructor(
99 > @ILogService logService: ILogService, byokLmProxyService.ts ×4
100 > @IByokLmBridgeRegistry private readonly _bridgeRegistry: IByokLmBridgeRegistry,
101 > ) {
102 > super(PROXY_USER_FACING_NAME, logService);
103 > }
105 > protected createState(): ByokLmProxyState {
106 > // No per-bind state — the bridge is resolved from the registry per request. byokLmProxyService.ts ×4
107 > return undefined;
108 > }
110 > async start(): Promise<IByokLmProxyHandle> {
111 > const { runtime, release } = await this.acquire(); byokLmProxyService.ts ×4
112 >
113 > let disposed = false;
114 > return {
115 > baseUrl: runtime.baseUrl,
116 > nonce: runtime.nonce,
117 > providerBaseUrl: (vendor: string) => `${runtime.baseUrl}${VENDOR_PATH_PREFIX}${encodeURIComponent(vendor)}`,
118 > dispose: () => {
119 > if (disposed) {
120 return;
121 }
122 > disposed = true; byokLmProxyService.ts ×4
123 > release();
124 > },
125 > };
126 > }
128 > /** Emit the base's fallback failure using the OpenAI error envelope. */
129 > protected override writeInternalError(res: http.ServerResponse): void {
130 this._writeJsonError(res, 500, 'Internal proxy error');
131 }
133 > protected override async handleRequest(req: http.IncomingMessage, res: http.ServerResponse, runtime: ILoopbackProxyRuntime<ByokLmProxyState>): Promise<void> {
134 > const method = req.method ?? 'GET'; byokLmProxyService.ts ×4
135 > const pathname = new URL(req.url ?? '/', 'http://127.0.0.1').pathname;
136 > this._logService.trace(`[${PROXY_USER_FACING_NAME}] ${method} ${pathname}`);
137 >
138 > if (method === 'GET' && pathname === '/') {
139 > res.writeHead(200, { 'Content-Type': 'text/plain' }); byokLmProxyService.ts ×1
140 > res.end('ok');
141 > return;
142 > }
144 > // Inbound requests carry `Bearer <nonce>.<sessionId>`; the runtime is
145 > // handed `<nonce>.<sessionId>` at session launch.
146 > const auth = parseProxyBearer(req.headers, runtime.nonce);
147 > if (!auth.valid || !auth.sessionId) { byokLmProxyService.ts ×4
148 > this._writeJsonError(res, 401, 'Invalid authentication', 'authentication_error'); byokLmProxyService.ts ×1
149 > return;
150 > }
152 > const vendor = this._parseVendorFromChatPath(pathname);
153 > if (method === 'POST' && vendor !== undefined) { byokLmProxyService.ts ×4
154 > await this._handleChatCompletions(req, res, runtime, vendor); byokLmProxyService.ts ×5
155 > return;
156 > }
158 > this._writeJsonError(res, 404, `No route for ${method} ${pathname}`, 'not_found_error');
161 > /**
162 > * Extract the vendor from a `/v/<vendor>/chat/completions` path, or return
163 > * `undefined` when the path is not a chat-completions route.
164 > */
165 > private _parseVendorFromChatPath(pathname: string): string | undefined {
166 > if (!pathname.startsWith(VENDOR_PATH_PREFIX) || !pathname.endsWith(CHAT_COMPLETIONS_SUFFIX)) { byokLmProxyService.ts ×4
167 > return undefined; byokLmProxyService.ts ×1
168 > }
169 > const vendorSegment = pathname.slice(VENDOR_PATH_PREFIX.length, pathname.length - CHAT_COMPLETIONS_SUFFIX.length); byokLmProxyService.ts ×3
170 > if (!vendorSegment) {
171 return undefined;
172 }
173 > let vendor: string; byokLmProxyService.ts ×3
174 > try {
175 > vendor = decodeURIComponent(vendorSegment);
176 > } catch {
177 return undefined;
178 }
179 > // Re-check for a path separator *after* decoding: a `%2F` survives the byokLmProxyService.ts ×3
180 > // pre-decode prefix/suffix checks but would decode into a second path
181 > // segment, breaking the single-segment `vendor/id` selection-id convention.
182 > if (!vendor || vendor.includes('/')) { byokLmProxyService.ts ×4
183 > return undefined; byokLmProxyService.ts ×1
184 > }
185 > return vendor; byokLmProxyService.ts ×5
188 > private async _handleChatCompletions(req: http.IncomingMessage, res: http.ServerResponse, runtime: ILoopbackProxyRuntime<ByokLmProxyState>, vendor: string): Promise<void> {
189 > let body: IOpenAiChatRequest; byokLmProxyService.ts ×5
190 > try {
191 > const raw = await readProxyRequestBody(req);
192 > body = JSON.parse(raw) as IOpenAiChatRequest;
193 > } catch (err) {
194 > this._writeJsonError(res, 400, `Invalid request body: ${err instanceof Error ? err.message : String(err)}`, 'invalid_request_error'); byokLmProxyService.ts ×1
195 > return;
196 > }
198 > let bridgeRequest;
199 > try {
200 > bridgeRequest = openAiRequestToBridge(vendor, body);
201 > } catch (err) {
202 const message = err instanceof OpenAiTranslationError ? err.message : String(err);
203 this._writeJsonError(res, 400, message, 'invalid_request_error');
204 return;
205 }
207 > const connection = this._bridgeRegistry.getServingConnection();
208 > if (!connection) {
209 > this._writeJsonError(res, 503, 'No renderer connection available to service BYOK models', 'api_error'); byokLmProxyService.ts ×1
210 > return;
211 > }
213 > // Register the request so {@link LoopbackProxyServer} aborts it on
214 > // teardown; a client-side disconnect also flips `clientGone` and aborts.
215 > // Both surface through the shared `AbortController`, which we re-check
216 > // after the async bridge hop before touching the response.
217 > const entry: IProxyInFlight = { ac: new AbortController(), res, clientGone: false };
218 > runtime.inFlight.add(entry);
219 > const onClose = () => {
220 entry.clientGone = true;
221 entry.ac.abort();
222 };
223 > res.on('close', onClose); byokLmProxyService.ts ×3
224 >
225 > try {
226 > const result = await connection.chat(bridgeRequest);
227 > if (entry.ac.signal.aborted || res.writableEnded) { byokLmProxyService.ts ×5
228 return;
229 }
230 > if (result.error) { byokLmProxyService.ts ×1
231 > this._writeJsonError(res, 502, result.error, 'api_error'); byokLmProxyService.ts ×1
232 > return;
233 > }
234 > res.writeHead(200, { byokLmProxyService.ts ×1
235 > 'Content-Type': 'text/event-stream',
236 > 'Cache-Control': 'no-cache',
237 > 'Connection': 'keep-alive',
238 > });
239 > for (const frame of bridgeResultToSseFrames(result, bridgeRequest.modelId)) {
240 > res.write(frame);
241 > }
242 > res.end();
243 > } catch (err) { byokLmProxyService.ts ×1
244 > if (entry.ac.signal.aborted || res.writableEnded) { byokLmProxyService.ts ×2
245 return;
246 }
247 > const message = err instanceof Error ? err.message : String(err); byokLmProxyService.ts ×2
248 > if (!res.headersSent) {
249 > this._writeJsonError(res, 502, message, 'api_error');
250 > } else {
251 try { res.end(); } catch { /* ignore */ }
252 }
253 > } finally { byokLmProxyService.ts ×3
254 > res.removeListener('close', onClose);
255 > runtime.inFlight.delete(entry);
256 > }
259 > private _writeJsonError(res: http.ServerResponse, status: number, message: string, type = 'api_error'): void {
260 > if (res.headersSent || res.writableEnded) { byokLmProxyService.ts ×2
261 return;
262 }
263 > res.writeHead(status, { 'Content-Type': 'application/json' }); byokLmProxyService.ts ×2
264 > res.end(openAiErrorBody(message, type));
265 > }
267 >
268 > /**
269 > * No-op {@link IByokLmProxyService} for agent host entrypoints that do not
270 > * support BYOK — e.g. the remote agent host, where no extension host runs
271 > * alongside the agent host to serve the renderer LM API.
272 > *
273 > */
274 > export class NullByokLmProxyService implements IByokLmProxyService {
275 >
276 > declare readonly _serviceBrand: undefined;
277 >
278 > start(): Promise<IByokLmProxyHandle> {
279 return Promise.reject(new Error('BYOK is not supported in this agent host'));
280 }
282 > dispose(): void {
283 // No-op: the null proxy never binds a socket, so there is nothing to close.
284 }