diff --git a/scripts/code-agent-host.js b/scripts/code-agent-host.js index 221e0ce60bf..ecd162a0a24 100644 --- a/scripts/code-agent-host.js +++ b/scripts/code-agent-host.js @@ -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 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 Log level to use (trace, debug, info, warning, error, off)\n' + - ' --help Show this help message', + ' --port Port to listen on (default: 8081, or VSCODE_AGENT_HOST_PORT env)\n' + + ' --connection-token A secret that must be included with all requests\n' + + ' --connection-token-file 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 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}`); diff --git a/src/vs/platform/agentHost/common/remoteAgentHostService.ts b/src/vs/platform/agentHost/common/remoteAgentHostService.ts index 375ece88206..1bc6ecb2a1d 100644 --- a/src/vs/platform/agentHost/common/remoteAgentHostService.ts +++ b/src/vs/platform/agentHost/common/remoteAgentHostService.ts @@ -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('remoteAgentHostService'); diff --git a/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts b/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts index 85ce773c838..45603675b65 100644 --- a/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts @@ -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())); } diff --git a/src/vs/platform/agentHost/electron-browser/remoteAgentHostServiceImpl.ts b/src/vs/platform/agentHost/electron-browser/remoteAgentHostServiceImpl.ts index c267f2765e2..085c7f2e965 100644 --- a/src/vs/platform/agentHost/electron-browser/remoteAgentHostServiceImpl.ts +++ b/src/vs/platform/agentHost/electron-browser/remoteAgentHostServiceImpl.ts @@ -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); diff --git a/src/vs/platform/agentHost/electron-browser/webSocketClientTransport.ts b/src/vs/platform/agentHost/electron-browser/webSocketClientTransport.ts index 6c32f0deccd..9b05e2b193c 100644 --- a/src/vs/platform/agentHost/electron-browser/webSocketClientTransport.ts +++ b/src/vs/platform/agentHost/electron-browser/webSocketClientTransport.ts @@ -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; diff --git a/src/vs/platform/agentHost/node/agentHostServerMain.ts b/src/vs/platform/agentHost/node/agentHostServerMain.ts index be7a5e2fb6b..dcae29e9eee 100644 --- a/src/vs/platform/agentHost/node/agentHostServerMain.ts +++ b/src/vs/platform/agentHost/node/agentHostServerMain.ts @@ -4,7 +4,7 @@ *--------------------------------------------------------------------------------------------*/ // Standalone agent host server with WebSocket protocol transport. -// Start with: node out/vs/platform/agentHost/node/agentHostServerMain.js [--port ] [--enable-mock-agent] [--quiet] [--log ] +// Start with: node out/vs/platform/agentHost/node/agentHostServerMain.js [--port ] [--connection-token ] [--connection-token-file ] [--without-connection-token] [--enable-mock-agent] [--quiet] [--log ] 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 { } // 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); } diff --git a/src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts index a6a3acf48d0..03361157689 100644 --- a/src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts @@ -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'], }); diff --git a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts index c2ad8b20cfc..b02f0a4d738 100644 --- a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts @@ -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'], },