src/vs/platform/agentHost/node/byokLmBridgeRegistry.ts
193 LOC · 183 covered · 10 uncovered · 35 ranges · 954 concepts · 18 introducers · 441 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
>
/*---------------------------------------------------------------------------------------------
byokLmBridgeRegistry.ts ×13
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 { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../base/common/lifecycle.js';
7
>
import { createDecorator } from '../../instantiation/common/instantiation.js';
8
>
import { IByokLmBridgeConnection, IByokLmModelInfo } from '../common/agentHostByokLm.js';
9
>
10
>
export const IByokLmBridgeRegistry = createDecorator<IByokLmBridgeRegistry>('byokLmBridgeRegistry');
11
>
12
>
/**
13
>
* Node-side registry of renderer {@link IByokLmBridgeConnection}s keyed by
14
>
* client id. Populated by the agent host's connection lifecycle (one entry per
15
>
* connected renderer) and consumed by {@link IByokLmProxyService} (inference
16
>
* routing) and {@link CopilotAgent} (model catalogue).
17
>
*
18
>
* **Single serving window, multiple connections.** BYOK is serviced by the
19
>
* renderer LM API, whose BYOK models are a property of the user's installed
20
>
* extensions, not of a particular window — so every window that registers the
21
>
* handler exposes the same set. Both the main workbench and the dedicated Agents
22
>
* app register it (each runs a full extension host whose LM API holds the same
23
>
* BYOK models), so either can serve. A connection that connects without binding
24
>
* the handler never pushes a snapshot and is treated as non-serving. The registry
25
>
* therefore does NOT aggregate per-window model sets; it surfaces the models from
26
>
* any one *serving* window (preferring one that actually has models) and routes
27
>
* inference there, automatically excluding non-serving windows.
28
>
*
29
>
* **Push, not pull.** Each connection pushes its current model snapshot over
30
>
* {@link IByokLmBridgeConnection.onDidChangeModels} (on subscribe and on every
31
>
* change); the registry subscribes on {@link register}, caches each snapshot, and
32
>
* fires {@link onDidChangeModels} when the serving model set changes. A connection
33
>
* becomes "serving" once it pushes its first snapshot (even an empty one).
34
>
*/
35
>
export interface IByokLmBridgeRegistry {
36
>
readonly _serviceBrand: undefined;
37
>
38
>
/** Register a renderer connection. Disposing the result removes it. */
39
>
register(clientId: string, connection: IByokLmBridgeConnection): IDisposable;
40
>
41
>
/**
42
>
* The serving window's BYOK models, read synchronously from the cache (no
43
>
* enumeration). Use this for fast reads driven by {@link onDidChangeModels}.
44
>
*/
45
>
getModels(): readonly IByokLmModelInfo[];
46
>
47
>
/**
48
>
* A connection that can serve BYOK inference, or `undefined` when none can.
49
>
* All serving windows expose the same models, so any one is a valid target.
50
>
*/
51
>
getServingConnection(): IByokLmBridgeConnection | undefined;
52
>
53
>
/**
54
>
* Subscribe to changes in the set of registered connections (a renderer
55
>
* connecting or disconnecting) or in the serving window's pushed models, so
56
>
* consumers can re-read {@link getModels}. Disposing the result removes the
57
>
* listener.
58
>
*/
59
>
onDidChangeModels(listener: () => void): IDisposable;
60
>
}
61
>
62
>
/**
63
>
* Per-connection registry entry. `models` is `undefined` until the connection
64
>
* pushes its first snapshot; a connection with defined `models` is "serving"
65
>
* (it pushed, even an empty list). Non-serving windows (those that did not
66
>
* register the BYOK handler and therefore never push) keep `models === undefined`.
67
>
*/
68
>
interface IConnectionEntry {
69
>
readonly connection: IByokLmBridgeConnection;
70
>
models: readonly IByokLmModelInfo[] | undefined;
71
>
readonly store: DisposableStore;
72
>
}
73
>
74
>
export class ByokLmBridgeRegistry implements IByokLmBridgeRegistry {
76
>
declare readonly _serviceBrand: undefined;
77
>
78
>
private readonly _entries = new Map<string, IConnectionEntry>();
79
>
private readonly _changeListeners = new Set<() => void>();
81
>
onDidChangeModels(listener: () => void): IDisposable {
83
>
return toDisposable(() => {
84
>
this._changeListeners.delete(listener);
85
>
});
86
>
}
88
>
private _notifyChanged(): void {
89
>
// Snapshot first: a listener may unsubscribe (mutating the set) while it
byokLmBridgeRegistry.ts ×4
90
>
// is being notified.
91
>
for (const listener of [...this._changeListeners]) {
93
>
}
96
>
register(clientId: string, connection: IByokLmBridgeConnection): IDisposable {
97
>
// Replace any prior entry for the same client id (e.g. a reconnect).
byokLmBridgeRegistry.ts ×4
98
>
this._entries.get(clientId)?.store.dispose();
99
>
100
>
const store = new DisposableStore();
101
>
const entry: IConnectionEntry = { connection, models: undefined, store };
102
>
this._entries.set(clientId, entry);
103
>
104
>
// Cache each pushed snapshot; notify only when the serving set changes.
105
>
store.add(connection.onDidChangeModels(models => {
107
>
if (this._entries.get(clientId) !== entry) {
108
return;
109
}
110
>
if (entry.models === undefined || !modelsEqual(entry.models, models)) {
byokLmBridgeRegistry.ts ×2
111
>
entry.models = models;
112
>
this._notifyChanged();
113
>
}
115
>
116
>
// The connection set changed (a renderer connected).
117
>
this._notifyChanged();
118
>
119
>
return toDisposable(() => {
120
>
if (this._entries.get(clientId) === entry) {
121
>
this._entries.delete(clientId);
122
>
entry.store.dispose();
123
>
this._notifyChanged();
124
>
}
125
>
});
126
>
}
128
>
getModels(): readonly IByokLmModelInfo[] {
130
>
}
132
>
getServingConnection(): IByokLmBridgeConnection | undefined {
134
>
}
136
>
/**
137
>
* A serving connection (`models` defined), preferring one whose model set is
138
>
* non-empty. All serving windows expose the same models, so any populated one
139
>
* is equivalent; the preference matters when a still-starting window pushes an
140
>
* empty list first — it must not shadow a peer that already has them. Falls
141
>
* back to a serving-but-empty window; non-serving windows are skipped.
142
>
*/
143
>
private _servingEntry(): IConnectionEntry | undefined {
145
>
for (const entry of this._entries.values()) {
148
>
}
151
>
}
153
>
}
157
>
158
>
/** Shallow structural comparison of two model lists (order-sensitive). */
159
>
function modelsEqual(a: readonly IByokLmModelInfo[], b: readonly IByokLmModelInfo[]): boolean {
byokLmBridgeRegistry.ts ×1
160
>
if (a.length !== b.length) {
162
>
}
164
>
const n = b[i];
165
>
return m.vendor === n.vendor && m.id === n.id && m.name === n.name && m.modelIdentifier === n.modelIdentifier && m.maxContextWindowTokens === n.maxContextWindowTokens && m.supportsVision === n.supportsVision;
166
>
});
167
>
}
169
>
/**
170
>
* No-op {@link IByokLmBridgeRegistry} for agent host entrypoints that do not
171
>
* support BYOK — e.g. the remote agent host, where no extension host runs
172
>
* alongside the agent host to serve the renderer LM API.
173
>
*/
174
>
export class NullByokLmBridgeRegistry implements IByokLmBridgeRegistry {
175
>
176
>
declare readonly _serviceBrand: undefined;
177
>
178
>
register(): IDisposable {
179
return Disposable.None;
180
}
182
>
getModels(): readonly IByokLmModelInfo[] {
183
return [];
184
}
186
>
getServingConnection(): IByokLmBridgeConnection | undefined {
187
return undefined;
188
}
190
>
onDidChangeModels(): IDisposable {
191
return Disposable.None;
192
}