Microsoft/vscode
1
0
Fork 0
mirror of https://github.com/microsoft/vscode.git synced 2026-04-02 00:09:30 +01:00
Code Issues Packages Projects Releases Wiki Activity

connection token support

Browse Source
This commit is contained in:
Connor Peet
2026-03-17 10:01:53 -07:00
parent f48e040708
commit f2f31ca7ea
8 changed files with 106 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
scripts/code-agent-host.js
View File

@@ -11,8 +11,8 @@ const minimist = require('minimist');
async function main() {
const args = minimist(process.argv.slice(2), {
boolean: ['help', 'enable-mock-agent', 'quiet'],
string: ['port', 'log'],
boolean: ['help', 'enable-mock-agent', 'quiet', 'without-connection-token'],
string: ['port', 'log', 'connection-token', 'connection-token-file'],
});
if (args.help) {
@@ -20,11 +20,14 @@ async function main() {
'Usage: ./scripts/code-agent-host.sh [options]\n' +
'\n' +
'Options:\n' +
' --port <number> Port to listen on (default: 8081, or VSCODE_AGENT_HOST_PORT env)\n' +
' --enable-mock-agent Enable the mock agent for testing\n' +
' --quiet Suppress logging output\n' +
' --log <level> Log level to use (trace, debug, info, warning, error, off)\n' +
' --help Show this help message',
' --port <number> Port to listen on (default: 8081, or VSCODE_AGENT_HOST_PORT env)\n' +
' --connection-token <token> A secret that must be included with all requests\n' +
' --connection-token-file <path> Path to a file containing the connection token\n' +
' --without-connection-token Run without a connection token\n' +
' --enable-mock-agent Enable the mock agent for testing\n' +
' --quiet Suppress logging output\n' +
' --log <level> Log level to use (trace, debug, info, warning, error, off)\n' +
' --help Show this help message',
);
return;
}
@@ -42,6 +45,15 @@ async function main() {
if (args.log) {
serverArgs.push('--log', String(args.log));
}
if (args['connection-token']) {
serverArgs.push('--connection-token', String(args['connection-token']));
}
if (args['connection-token-file']) {
serverArgs.push('--connection-token-file', String(args['connection-token-file']));
}
if (args['without-connection-token']) {
serverArgs.push('--without-connection-token');
}
const addr = await startServer(serverArgs);
console.log(`Agent Host server listening on ${addr}`);

1
src/vs/platform/agentHost/common/remoteAgentHostService.ts
View File

@@ -15,6 +15,7 @@ export const RemoteAgentHostsSettingId = 'chat.remoteAgentHosts';
export interface IRemoteAgentHostEntry {
readonly address: string;
readonly name: string;
readonly connectionToken?: string;
}
export const IRemoteAgentHostService = createDecorator<IRemoteAgentHostService>('remoteAgentHostService');

3
src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts
View File

@@ -66,10 +66,11 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC
constructor(
address: string,
connectionToken: string | undefined,
@ILogService private readonly _logService: ILogService,
) {
super();
this._transport = this._register(new WebSocketClientTransport(address));
this._transport = this._register(new WebSocketClientTransport(address, connectionToken));
this._register(this._transport.onMessage(msg => this._handleMessage(msg)));
this._register(this._transport.onClose(() => this._onDidClose.fire()));
}

6
src/vs/platform/agentHost/electron-browser/remoteAgentHostServiceImpl.ts
View File

@@ -114,7 +114,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo
// Add new connections
for (const entry of entries) {
if (!this._entries.has(entry.address)) {
this._connectTo(entry.address);
this._connectTo(entry.address, entry.connectionToken);
}
}
@@ -124,9 +124,9 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo
}
}
private _connectTo(address: string): void {
private _connectTo(address: string, connectionToken?: string): void {
const store = new DisposableStore();
const client = store.add(this._instantiationService.createInstance(RemoteAgentHostProtocolClient, address));
const client = store.add(this._instantiationService.createInstance(RemoteAgentHostProtocolClient, address, connectionToken));
const entry: IConnectionEntry = { store, client, connected: false };
this._entries.set(address, entry);

13
src/vs/platform/agentHost/electron-browser/webSocketClientTransport.ts
View File

@@ -8,6 +8,7 @@
import { Emitter } from '../../../base/common/event.js';
import { Disposable } from '../../../base/common/lifecycle.js';
import { connectionTokenQueryName } from '../../../base/common/network.js';
import type { IProtocolMessage } from '../common/state/sessionProtocol.js';
import type { IProtocolTransport } from '../common/state/sessionTransport.js';
import { protocolReplacer, protocolReviver } from '../common/state/jsonSerialization.js';
@@ -36,7 +37,10 @@ export class WebSocketClientTransport extends Disposable implements IProtocolTra
return this._ws?.readyState === WebSocket.OPEN;
}
constructor(private readonly _address: string) {
constructor(
private readonly _address: string,
private readonly _connectionToken?: string,
) {
super();
}
@@ -51,10 +55,15 @@ export class WebSocketClientTransport extends Disposable implements IProtocolTra
return;
}
const url = this._address.startsWith('ws://') || this._address.startsWith('wss://')
let url = this._address.startsWith('ws://') || this._address.startsWith('wss://')
? this._address
: `ws://${this._address}`;
if (this._connectionToken) {
const separator = url.includes('?') ? '&' : '?';
url += `${separator}${connectionTokenQueryName}=${encodeURIComponent(this._connectionToken)}`;
}
const ws = new WebSocket(url);
this._ws = ws;

77
src/vs/platform/agentHost/node/agentHostServerMain.ts
View File

@@ -4,7 +4,7 @@
*--------------------------------------------------------------------------------------------*/
// Standalone agent host server with WebSocket protocol transport.
// Start with: node out/vs/platform/agentHost/node/agentHostServerMain.js [--port <port>] [--enable-mock-agent] [--quiet] [--log <level>]
// Start with: node out/vs/platform/agentHost/node/agentHostServerMain.js [--port <port>] [--connection-token <token>] [--connection-token-file <path>] [--without-connection-token] [--enable-mock-agent] [--quiet] [--log <level>]
import { fileURLToPath } from 'url';
@@ -13,8 +13,10 @@ import { fileURLToPath } from 'url';
// This file lives at out/vs/platform/agentHost/node/ - the root is `out/`.
globalThis._VSCODE_FILE_ROOT = fileURLToPath(new URL('../../../..', import.meta.url));
import * as fs from 'fs';
import { DisposableStore } from '../../../base/common/lifecycle.js';
import { observableValue } from '../../../base/common/observable.js';
import { generateUuid } from '../../../base/common/uuid.js';
import { localize } from '../../../nls.js';
import { NativeEnvironmentService } from '../../environment/node/environmentService.js';
import { INativeEnvironmentService } from '../../environment/common/environment.js';
@@ -40,10 +42,14 @@ function log(msg: string): void {
// ---- Options ----------------------------------------------------------------
const connectionTokenRegex = /^[0-9A-Za-z_-]+$/;
interface IServerOptions {
readonly port: number;
readonly enableMockAgent: boolean;
readonly quiet: boolean;
/** Connection token string, or `undefined` when `--without-connection-token`. */
readonly connectionToken: string | undefined;
}
function parseServerOptions(): IServerOptions {
@@ -53,7 +59,48 @@ function parseServerOptions(): IServerOptions {
const port = portIdx >= 0 ? parseInt(argv[portIdx + 1], 10) : envPort;
const enableMockAgent = argv.includes('--enable-mock-agent');
const quiet = argv.includes('--quiet');
return { port, enableMockAgent, quiet };
// Connection token
const withoutConnectionToken = argv.includes('--without-connection-token');
const connectionTokenIdx = argv.indexOf('--connection-token');
const connectionTokenFileIdx = argv.indexOf('--connection-token-file');
const rawToken = connectionTokenIdx >= 0 ? argv[connectionTokenIdx + 1] : undefined;
const tokenFilePath = connectionTokenFileIdx >= 0 ? argv[connectionTokenFileIdx + 1] : undefined;
let connectionToken: string | undefined;
if (withoutConnectionToken) {
if (rawToken !== undefined || tokenFilePath !== undefined) {
log('Error: --without-connection-token cannot be used with --connection-token or --connection-token-file');
process.exit(1);
}
connectionToken = undefined;
} else if (tokenFilePath !== undefined) {
if (rawToken !== undefined) {
log('Error: --connection-token cannot be used with --connection-token-file');
process.exit(1);
}
try {
connectionToken = fs.readFileSync(tokenFilePath).toString().replace(/\r?\n$/, '');
} catch {
log(`Error: Unable to read connection token file at '${tokenFilePath}'`);
process.exit(1);
}
if (!connectionTokenRegex.test(connectionToken!)) {
log(`Error: The connection token in '${tokenFilePath}' does not adhere to the characters 0-9, a-z, A-Z, _, or -.`);
process.exit(1);
}
} else if (rawToken !== undefined) {
if (!connectionTokenRegex.test(rawToken)) {
log(`Error: The connection token '${rawToken}' does not adhere to the characters 0-9, a-z, A-Z, _, or -.`);
process.exit(1);
}
connectionToken = rawToken;
} else {
// Default: generate a random token (secure by default)
connectionToken = generateUuid();
}
return { port, enableMockAgent, quiet, connectionToken };
}
// ---- Main -------------------------------------------------------------------
@@ -136,25 +183,37 @@ async function main(): Promise<void> {
}
// WebSocket server
const wsServer = disposables.add(await WebSocketProtocolServer.create(options.port, logService));
const wsServer = disposables.add(await WebSocketProtocolServer.create({
port: options.port,
connectionTokenValidate: options.connectionToken
? token => token === options.connectionToken
: undefined,
}, logService));
// Wire up protocol handler
disposables.add(new ProtocolServerHandler(stateManager, wsServer, sideEffects, logService));
// Report ready
function reportReady(addr: string): void {
const listeningPort = addr.split(':').pop();
let wsUrl = `ws://${addr}`;
if (options.connectionToken) {
wsUrl += `?tkn=${options.connectionToken}`;
}
process.stdout.write(`READY:${listeningPort}\n`);
log(`WebSocket server listening on ${wsUrl}`);
logService.info(`[AgentHostServer] WebSocket server listening on ${wsUrl}`);
}
const address = wsServer.address;
if (address) {
const listeningPort = address.split(':').pop();
process.stdout.write(`READY:${listeningPort}\n`);
logService.info(`[AgentHostServer] WebSocket server listening on ws://${address}`);
reportReady(address);
} else {
const interval = setInterval(() => {
const addr = wsServer.address;
if (addr) {
clearInterval(interval);
const listeningPort = addr.split(':').pop();
process.stdout.write(`READY:${listeningPort}\n`);
logService.info(`[AgentHostServer] WebSocket server listening on ws://${addr}`);
reportReady(addr);
}
}, 10);
}

2
src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts
View File

@@ -189,7 +189,7 @@ class TestProtocolClient {
async function startServer(): Promise<{ process: ChildProcess; port: number }> {
return new Promise((resolve, reject) => {
const serverPath = fileURLToPath(new URL('../../node/agentHostServerMain.js', import.meta.url));
const child = fork(serverPath, ['--enable-mock-agent', '--quiet', '--port', '0'], {
const child = fork(serverPath, ['--enable-mock-agent', '--quiet', '--port', '0', '--without-connection-token'], {
stdio: ['pipe', 'pipe', 'pipe', 'ipc'],
});

1
src/vs/workbench/contrib/chat/browser/chat.contribution.ts
View File

@@ -723,6 +723,7 @@ configurationRegistry.registerConfiguration({
properties: {
address: { type: 'string', description: nls.localize('chat.remoteAgentHosts.address', "The address of the remote agent host (e.g. \"localhost:3000\").") },
name: { type: 'string', description: nls.localize('chat.remoteAgentHosts.name', "A display name for this remote agent host.") },
connectionToken: { type: 'string', description: nls.localize('chat.remoteAgentHosts.connectionToken', "An optional connection token for authenticating with the remote agent host.") },
},
required: ['address', 'name'],
},