diff --git a/packages/agent-core-v2/package.json b/packages/agent-core-v2/package.json index 00082474cd..5795b574a6 100644 --- a/packages/agent-core-v2/package.json +++ b/packages/agent-core-v2/package.json @@ -85,6 +85,7 @@ "tar": "^7.5.13", "ulid": "^3.0.1", "undici": "^7.27.1", + "ws": "^8.18.0", "yauzl": "^3.3.0", "yazl": "^3.3.1", "zod": "^4.3.6" @@ -98,6 +99,7 @@ "@types/retry": "0.12.0", "@types/sinon": "^21.0.1", "@types/tar": "^7.0.87", + "@types/ws": "^8.18.0", "@types/yauzl": "^2.10.3", "@types/yazl": "^2.4.6", "@vitejs/plugin-react": "^4.4.1", diff --git a/packages/agent-core-v2/src/app/remoteControl/flag.ts b/packages/agent-core-v2/src/app/remoteControl/flag.ts new file mode 100644 index 0000000000..8ce970c1d8 --- /dev/null +++ b/packages/agent-core-v2/src/app/remoteControl/flag.ts @@ -0,0 +1,19 @@ +/** + * `remoteControl` domain — experimental device tunnel gate. + */ + +import { type FlagDefinitionInput, registerFlagDefinition } from '#/app/flag/flagRegistry'; + +export const REMOTE_CONTROL_FLAG_ID = 'remote_control'; +export const REMOTE_CONTROL_FLAG_ENV = 'KIMI_CODE_EXPERIMENTAL_REMOTE_CONTROL'; + +export const remoteControlFlag: FlagDefinitionInput = { + id: REMOTE_CONTROL_FLAG_ID, + title: 'Remote Control', + description: 'Allow an authenticated relay to proxy this local Kimi Code server.', + env: REMOTE_CONTROL_FLAG_ENV, + default: false, + surface: 'core', +}; + +registerFlagDefinition(remoteControlFlag); diff --git a/packages/agent-core-v2/src/app/remoteControl/protocol.ts b/packages/agent-core-v2/src/app/remoteControl/protocol.ts new file mode 100644 index 0000000000..b71ec962a6 --- /dev/null +++ b/packages/agent-core-v2/src/app/remoteControl/protocol.ts @@ -0,0 +1,235 @@ +/** + * `remoteControl` domain — validated relay protocol and HTTP tunnel codecs. + */ + +import { z } from 'zod'; + +const HeadersSchema = z.record(z.string(), z.string()); + +export const RemoteDisconnectReasonSchema = z.enum([ + 'user_requested', + 'server_shutting_down', + 'local_server_stopped', + 'client_upgrading', +]); +export type RemoteDisconnectReason = z.infer; + +export const RemoteStreamCloseReasonSchema = z.enum([ + 'browser_closed', + 'timeout', + 'server_shutting_down', + 'local_closed', +]); +export type RemoteStreamCloseReason = z.infer; + +export const RemoteStreamErrorCodeSchema = z.enum([ + 'LOCAL_WS_FAILED', + 'TUNNEL_STREAM_FAILED', + 'TIMEOUT', + 'UNKNOWN', +]); +export type RemoteStreamErrorCode = z.infer; + +export const ManagementInboundMessageSchema = z.discriminatedUnion('type', [ + z.object({ + type: z.literal('register_ack'), + payload: z.object({ success: z.boolean() }), + }), + z.object({ + type: z.literal('open_ws'), + payload: z.object({ + stream_id: z.string().min(1), + path: z.string().min(1), + headers: HeadersSchema.default({}), + }), + }), + z.object({ + type: z.literal('close_ws'), + payload: z.object({ + stream_id: z.string().min(1), + close_code: z.number().int().min(1000).max(4999).optional(), + reason: RemoteStreamCloseReasonSchema, + path: z.string().min(1).optional(), + headers: HeadersSchema.optional(), + }), + }), + z.object({ + type: z.literal('disconnect'), + payload: z.object({ reason: RemoteDisconnectReasonSchema }), + }), +]); +export type ManagementInboundMessage = z.infer; + +export const ManagementOutboundMessageSchema = z.discriminatedUnion('type', [ + z.object({ + type: z.literal('register'), + payload: z.object({ + device_id: z.string().min(1), + alias: z.string().min(1), + platform: z.string().min(1), + client_version: z.string().min(1), + local_base_url: z.string().url(), + }), + }), + z.object({ + type: z.literal('open_ws_result'), + payload: z.object({ + stream_id: z.string().min(1), + success: z.boolean(), + error_code: RemoteStreamErrorCodeSchema.optional(), + error_message: z.string().optional(), + }), + }), + z.object({ + type: z.literal('disconnect'), + payload: z.object({ reason: RemoteDisconnectReasonSchema }), + }), +]); +export type ManagementOutboundMessage = z.infer; + +export const HttpTunnelMessageSchema = z.object({ + request_id: z.string().min(1), + type: z.enum(['request', 'response', 'response_chunk']), + is_last: z.boolean(), + body_base64: z.string(), +}); +export type HttpTunnelMessage = z.infer; + +export interface LocalHttpRequest { + readonly method: string; + readonly path: string; + readonly headers: Readonly>; + readonly body: Buffer; +} + +export function parseManagementInbound(data: string | Buffer): ManagementInboundMessage { + return ManagementInboundMessageSchema.parse(JSON.parse(data.toString())); +} + +export function parseHttpTunnelMessage(data: string | Buffer): HttpTunnelMessage { + return HttpTunnelMessageSchema.parse(JSON.parse(data.toString())); +} + +export function encodeManagementMessage(message: ManagementOutboundMessage): string { + return JSON.stringify(ManagementOutboundMessageSchema.parse(message)); +} + +export function encodeHttpTunnelMessage(message: HttpTunnelMessage): string { + return JSON.stringify(HttpTunnelMessageSchema.parse(message)); +} + +export class HttpRequestAssembler { + private readonly chunks = new Map(); + + push(message: HttpTunnelMessage): LocalHttpRequest | undefined { + if (message.type !== 'request') throw new Error('expected an HTTP tunnel request'); + const chunks = this.chunks.get(message.request_id) ?? []; + chunks.push(Buffer.from(message.body_base64, 'base64')); + if (!message.is_last) { + this.chunks.set(message.request_id, chunks); + return undefined; + } + this.chunks.delete(message.request_id); + return parseRawHttpRequest(Buffer.concat(chunks)); + } + + clear(): void { + this.chunks.clear(); + } +} + +const HOP_BY_HOP_HEADERS = new Set([ + 'connection', + 'keep-alive', + 'proxy-authenticate', + 'proxy-authorization', + 'te', + 'trailer', + 'transfer-encoding', + 'upgrade', +]); + +export function sanitizeForwardHeaders( + headers: Readonly>, +): Record { + const connectionTokens = new Set( + Object.entries(headers) + .find(([name]) => name.toLowerCase() === 'connection')?.[1] + .split(',') + .map((value) => value.trim().toLowerCase()) + .filter(Boolean) ?? [], + ); + const result: Record = {}; + for (const [name, value] of Object.entries(headers)) { + const lower = name.toLowerCase(); + if ( + HOP_BY_HOP_HEADERS.has(lower) || + connectionTokens.has(lower) || + lower === 'authorization' || + lower === 'proxy-authorization' || + lower === 'host' || + lower === 'content-length' || + lower === 'origin' || + lower.startsWith('sec-websocket-') + ) continue; + result[name] = value; + } + return result; +} + +export function resolveLocalUrl(localBaseUrl: string, path: string): URL { + if (!path.startsWith('/')) throw new Error('local proxy path must start with /'); + const base = new URL(localBaseUrl); + if (base.protocol !== 'http:' && base.protocol !== 'https:') { + throw new Error('local base URL must use HTTP or HTTPS'); + } + const url = new URL(path, base); + if (url.origin !== base.origin) throw new Error('local proxy path must stay on the local origin'); + return url; +} + +export function parseRawHttpRequest(raw: Buffer): LocalHttpRequest { + const separator = raw.indexOf('\r\n\r\n'); + if (separator < 0) throw new Error('HTTP request headers are incomplete'); + const lines = raw.subarray(0, separator).toString('latin1').split('\r\n'); + const requestLine = lines.shift(); + const match = requestLine?.match(/^([A-Z]+)\s+(\S+)\s+HTTP\/1\.[01]$/); + if (match === null || match === undefined) throw new Error('invalid HTTP request line'); + const headers: Record = {}; + for (const line of lines) { + const colon = line.indexOf(':'); + if (colon <= 0) throw new Error('invalid HTTP request header'); + const name = line.slice(0, colon).trim(); + const value = line.slice(colon + 1).trim(); + headers[name] = headers[name] === undefined ? value : `${headers[name]}, ${value}`; + } + return { + method: match[1]!, + path: match[2]!, + headers: sanitizeForwardHeaders(headers), + body: raw.subarray(separator + 4), + }; +} + +export function serializeHttpResponseHead( + statusCode: number, + statusMessage: string, + headers: Readonly>, + bodyLength?: number, +): Buffer { + const safeHeaders: Record = {}; + for (const [name, value] of Object.entries(headers)) { + const lower = name.toLowerCase(); + if (value === undefined || HOP_BY_HOP_HEADERS.has(lower) || lower === 'content-length') continue; + safeHeaders[name] = typeof value === 'string' ? value : value.join(', '); + } + if (bodyLength !== undefined) safeHeaders['Content-Length'] = String(bodyLength); + const lines = [`HTTP/1.1 ${String(statusCode)} ${statusMessage}`]; + for (const [name, value] of Object.entries(safeHeaders)) lines.push(`${name}: ${value}`); + return Buffer.from(`${lines.join('\r\n')}\r\n\r\n`, 'latin1'); +} + +export const BAD_GATEWAY_RESPONSE = Buffer.from( + 'HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n', + 'latin1', +); diff --git a/packages/agent-core-v2/src/app/remoteControl/remoteControl.ts b/packages/agent-core-v2/src/app/remoteControl/remoteControl.ts new file mode 100644 index 0000000000..03769d4f4a --- /dev/null +++ b/packages/agent-core-v2/src/app/remoteControl/remoteControl.ts @@ -0,0 +1,30 @@ +/** + * `remoteControl` domain — App-scope device tunnel lifecycle contract. + */ + +import type { Event } from '#/_base/event'; +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; + +import type { RemoteDisconnectReason } from './protocol'; + +export type RemoteControlState = 'disabled' | 'offline' | 'connecting' | 'online'; + +export interface RemoteControlStartOptions { + readonly relayBaseUrl: string; + readonly localBaseUrl: string; + readonly alias: string; + readonly getLocalToken: () => string; + /** Stream setup deadline in milliseconds. Defaults to 10 seconds; tests may lower it. */ + readonly streamOpenTimeoutMs?: number; +} + +export interface IRemoteControlService { + readonly _serviceBrand: undefined; + readonly state: RemoteControlState; + readonly onDidChangeState: Event; + start(options: RemoteControlStartOptions): Promise; + stop(reason: Extract): Promise; +} + +export const IRemoteControlService: ServiceIdentifier = + createDecorator('remoteControlService'); diff --git a/packages/agent-core-v2/src/app/remoteControl/remoteControlService.ts b/packages/agent-core-v2/src/app/remoteControl/remoteControlService.ts new file mode 100644 index 0000000000..e847f8cf51 --- /dev/null +++ b/packages/agent-core-v2/src/app/remoteControl/remoteControlService.ts @@ -0,0 +1,519 @@ +/** + * `remoteControl` domain — `IRemoteControlService` implementation. + * + * Orchestrates OAuth-backed relay connections through `auth`, reads stable host + * identity through `bootstrap`, checks its experimental gate through `flag`, and + * records credential-free lifecycle diagnostics through `log`. Bound at App scope. + */ + +import { createKimiDeviceId, KIMI_CODE_PROVIDER_NAME } from '@moonshot-ai/kimi-code-oauth'; + +import { Disposable, DisposableStore, type IDisposable } from '#/_base/di/lifecycle'; +import { LifecycleScope, ScopeActivation, registerScopedService } from '#/_base/di/scope'; +import { Emitter, type Event } from '#/_base/event'; +import { ILogService } from '#/_base/log/log'; +import { IOAuthService } from '#/app/auth/auth'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { IFlagService } from '#/app/flag/flag'; + +import { REMOTE_CONTROL_FLAG_ID } from './flag'; +import { + BAD_GATEWAY_RESPONSE, + HttpRequestAssembler, + encodeHttpTunnelMessage, + encodeManagementMessage, + parseHttpTunnelMessage, + parseManagementInbound, + resolveLocalUrl, + serializeHttpResponseHead, + type HttpTunnelMessage, + type LocalHttpRequest, + type ManagementInboundMessage, + type RemoteDisconnectReason, + type RemoteStreamErrorCode, +} from './protocol'; +import { + IRemoteControlService, + type RemoteControlStartOptions, + type RemoteControlState, +} from './remoteControl'; +import { + IRemoteControlTransport, + type RemoteSocket, + type RemoteSocketBridge, +} from './remoteControlTransport'; + +const STREAM_OPEN_TIMEOUT_MS = 10_000; +const MAX_RECONNECT_DELAY_MS = 30_000; + +interface StreamHandle { + readonly bridge: RemoteSocketBridge; + closeListener?: IDisposable; +} + +export class RemoteControlService extends Disposable implements IRemoteControlService { + declare readonly _serviceBrand: undefined; + + private readonly _onDidChangeState = this._register(new Emitter()); + readonly onDidChangeState: Event = this._onDidChangeState.event; + private readonly connectionResources = this._register(new DisposableStore()); + private readonly streams = new Map(); + private readonly requestControllers = new Map(); + private httpTunnelOpening = false; + private readonly assembler = new HttpRequestAssembler(); + private options: RemoteControlStartOptions | undefined; + private management: RemoteSocket | undefined; + private httpTunnel: RemoteSocket | undefined; + private rootController: AbortController | undefined; + private reconnectTimer: ReturnType | undefined; + private reconnectAttempts = 0; + private generation = 0; + private stopping = true; + private _state: RemoteControlState = 'disabled'; + + get state(): RemoteControlState { + return this._state; + } + + constructor( + @IRemoteControlTransport private readonly transport: IRemoteControlTransport, + @IOAuthService private readonly oauth: IOAuthService, + @IBootstrapService private readonly bootstrap: IBootstrapService, + @IFlagService private readonly flags: IFlagService, + @ILogService private readonly log: ILogService, + ) { + super(); + } + + async start(options: RemoteControlStartOptions): Promise { + if (!this.flags.enabled(REMOTE_CONTROL_FLAG_ID)) { + this.setState('disabled'); + return; + } + validateStartOptions(options); + if (!this.stopping) await this.stop('user_requested'); + this.options = options; + this.stopping = false; + this.reconnectAttempts = 0; + await this.connect(); + } + + async stop( + reason: Extract, + ): Promise { + if (this.stopping) return; + this.stopping = true; + this.generation++; + this.clearReconnectTimer(); + if (this.management !== undefined) { + this.management.send(encodeManagementMessage({ type: 'disconnect', payload: { reason } })); + await Promise.resolve(); + } + this.resetConnections(); + this.setState('offline'); + } + + override dispose(): void { + this.stopping = true; + this.generation++; + this.clearReconnectTimer(); + this.resetConnections(); + super.dispose(); + } + + private async connect(): Promise { + const options = this.options; + if (options === undefined || this.stopping) return; + const generation = ++this.generation; + this.resetConnections(); + this.rootController = new AbortController(); + this.setState('connecting'); + + try { + const token = await this.oauth.getCachedAccessToken(KIMI_CODE_PROVIDER_NAME); + if (token === undefined || token.length === 0) { + this.setState('offline'); + this.log.info('remote control is waiting for OAuth authentication'); + this.scheduleReconnect(false); + return; + } + const deviceId = createKimiDeviceId(this.bootstrap.homeDir); + const management = await this.transport.connectManagement( + options.relayBaseUrl, + token, + this.rootController.signal, + ); + if (generation !== this.generation || this.stopping) { + management.dispose(); + return; + } + this.management = management; + this.bindManagement(management, token, generation, deviceId); + management.send(encodeManagementMessage({ + type: 'register', + payload: { + device_id: deviceId, + alias: options.alias, + platform: this.bootstrap.platform, + client_version: this.bootstrap.clientIdentity.version, + local_base_url: options.localBaseUrl, + }, + })); + } catch (error) { + if (generation !== this.generation || this.stopping) return; + this.setState('offline'); + this.log.warn('remote control connection failed', { error_type: errorType(error) }); + this.scheduleReconnect(false); + } + } + + private bindManagement(socket: RemoteSocket, token: string, generation: number, deviceId: string): void { + this.connectionResources.add(socket.onMessage((data, binary) => { + if (binary) { + this.protocolError(socket, generation); + return; + } + let message: ManagementInboundMessage; + try { + message = parseManagementInbound(data); + } catch { + this.protocolError(socket, generation); + return; + } + void this.handleManagementMessage(message, token, generation, deviceId); + })); + this.connectionResources.add(socket.onPing((data) => { socket.pong(data); })); + this.connectionResources.add(socket.onClose(() => { this.handleConnectionLoss(generation); })); + this.connectionResources.add(socket.onError(() => { this.handleConnectionLoss(generation); })); + } + + private async handleManagementMessage( + message: ManagementInboundMessage, + token: string, + generation: number, + deviceId: string, + ): Promise { + if (generation !== this.generation || this.stopping) return; + switch (message.type) { + case 'register_ack': + if (!message.payload.success) { + this.management?.close(1008, 'registration rejected'); + return; + } + await this.openHttpTunnel(token, generation, deviceId); + return; + case 'open_ws': + await this.openStream(message.payload, token, generation); + return; + case 'close_ws': + this.closeStream(message.payload.stream_id, message.payload.close_code, message.payload.reason); + if ( + message.payload.reason === 'server_shutting_down' && + message.payload.path !== undefined + ) { + await this.openStream({ + stream_id: message.payload.stream_id, + path: message.payload.path, + headers: message.payload.headers ?? {}, + }, token, generation); + } + return; + case 'disconnect': + if (message.payload.reason === 'user_requested') { + this.stopping = true; + this.generation++; + this.resetConnections(); + this.setState('offline'); + } else { + this.handleConnectionLoss(generation, true); + } + } + } + + private async openHttpTunnel(token: string, generation: number, deviceId: string): Promise { + if ( + this.options === undefined || + this.rootController === undefined || + this.httpTunnel !== undefined || + this.httpTunnelOpening + ) return; + this.httpTunnelOpening = true; + try { + const socket = await this.transport.connectHttpTunnel( + this.options.relayBaseUrl, + token, + deviceId, + this.rootController.signal, + ); + if (generation !== this.generation || this.stopping) { + socket.dispose(); + return; + } + this.httpTunnel = socket; + this.connectionResources.add(socket.onMessage((data, binary) => { + if (binary) { + this.protocolError(socket, generation); + return; + } + try { + const message = parseHttpTunnelMessage(data); + if (message.type !== 'request') { + this.protocolError(socket, generation); + return; + } + const request = this.assembler.push(message); + if (request !== undefined) void this.forwardHttp(message.request_id, request, generation); + } catch { + this.protocolError(socket, generation); + } + })); + this.connectionResources.add(socket.onPing((data) => { socket.pong(data); })); + this.connectionResources.add(socket.onClose(() => { this.handleConnectionLoss(generation); })); + this.connectionResources.add(socket.onError(() => { this.handleConnectionLoss(generation); })); + this.reconnectAttempts = 0; + this.setState('online'); + } catch (error) { + if (generation !== this.generation || this.stopping) return; + this.log.warn('remote control HTTP tunnel failed', { error_type: errorType(error) }); + this.handleConnectionLoss(generation); + } finally { + this.httpTunnelOpening = false; + } + } + + private async forwardHttp( + requestId: string, + request: LocalHttpRequest, + generation: number, + ): Promise { + const options = this.options; + const root = this.rootController; + if (options === undefined || root === undefined) return; + const controller = new AbortController(); + const abort = (): void => { controller.abort(root.signal.reason); }; + root.signal.addEventListener('abort', abort, { once: true }); + this.requestControllers.set(requestId, controller); + try { + const response = await this.transport.forwardLocalHttp( + options.localBaseUrl, + request, + options.getLocalToken(), + controller.signal, + ); + if (generation !== this.generation || controller.signal.aborted) return; + if (response.streaming) { + this.sendHttp(requestId, 'response_chunk', false, + serializeHttpResponseHead(response.statusCode, response.statusMessage, response.headers)); + for await (const chunk of response.body) { + if (controller.signal.aborted) return; + this.sendHttp(requestId, 'response_chunk', false, chunk); + } + this.sendHttp(requestId, 'response_chunk', true, Buffer.alloc(0)); + } else { + const chunks: Buffer[] = []; + for await (const chunk of response.body) chunks.push(chunk); + const body = Buffer.concat(chunks); + const head = serializeHttpResponseHead( + response.statusCode, + response.statusMessage, + response.headers, + body.byteLength, + ); + this.sendHttp(requestId, 'response', true, Buffer.concat([head, body])); + } + } catch { + if (!controller.signal.aborted) this.sendHttp(requestId, 'response', true, BAD_GATEWAY_RESPONSE); + } finally { + root.signal.removeEventListener('abort', abort); + this.requestControllers.delete(requestId); + } + } + + private sendHttp( + requestId: string, + type: Extract, + isLast: boolean, + body: Buffer, + ): void { + this.httpTunnel?.send(encodeHttpTunnelMessage({ + request_id: requestId, + type, + is_last: isLast, + body_base64: body.toString('base64'), + })); + } + + private async openStream( + payload: { stream_id: string; path: string; headers: Record }, + token: string, + generation: number, + ): Promise { + const options = this.options; + const root = this.rootController; + if (options === undefined || root === undefined) return; + this.closeStream(payload.stream_id); + const controller = new AbortController(); + const timeout = setTimeout( + () => { controller.abort(new Error('timeout')); }, + options.streamOpenTimeoutMs ?? STREAM_OPEN_TIMEOUT_MS, + ); + const abort = (): void => { controller.abort(root.signal.reason); }; + root.signal.addEventListener('abort', abort, { once: true }); + let local: RemoteSocket | undefined; + let tunnel: RemoteSocket | undefined; + // Open the relay stream before the local WebSocket: the local server emits + // its first frames (e.g. `server_hello`) immediately on connect, and any of + // them sent before the tunnel is up would be lost on the floor. + let stage: RemoteStreamErrorCode = 'TUNNEL_STREAM_FAILED'; + try { + tunnel = await this.transport.connectTunnelStream( + options.relayBaseUrl, + payload.stream_id, + token, + controller.signal, + ); + stage = 'LOCAL_WS_FAILED'; + local = await this.transport.connectLocalWebSocket( + options.localBaseUrl, + payload.path, + payload.headers, + options.getLocalToken(), + controller.signal, + ); + if (generation !== this.generation || controller.signal.aborted) throw controller.signal.reason; + stage = 'UNKNOWN'; + const bridge = this.transport.bridgeWebSockets(local, tunnel); + const handle: StreamHandle = { bridge }; + this.streams.set(payload.stream_id, handle); + handle.closeListener = bridge.onClose(() => { + if (this.streams.get(payload.stream_id)?.bridge === bridge) { + this.streams.delete(payload.stream_id); + } + }); + this.management?.send(encodeManagementMessage({ + type: 'open_ws_result', + payload: { stream_id: payload.stream_id, success: true }, + })); + } catch (error) { + local?.dispose(); + tunnel?.dispose(); + const errorCode = controller.signal.aborted && error instanceof Error && error.message === 'timeout' + ? 'TIMEOUT' + : stage; + this.management?.send(encodeManagementMessage({ + type: 'open_ws_result', + payload: { + stream_id: payload.stream_id, + success: false, + error_code: errorCode, + error_message: streamErrorMessage(errorCode), + }, + })); + } finally { + clearTimeout(timeout); + root.signal.removeEventListener('abort', abort); + } + } + + private closeStream(streamId: string, code = 1000, reason = ''): void { + const stream = this.streams.get(streamId); + if (stream === undefined) return; + this.streams.delete(streamId); + stream.closeListener?.dispose(); + stream.bridge.close(code, reason); + this.log.debug('remote control stream closed', { stream_id: streamId, close_code: code, reason }); + } + + private protocolError(socket: RemoteSocket, generation: number): void { + socket.close(1002, 'protocol error'); + this.handleConnectionLoss(generation); + } + + private handleConnectionLoss(generation: number, immediate = false): void { + if (generation !== this.generation || this.stopping) return; + this.generation++; + this.resetConnections(); + this.setState('offline'); + this.scheduleReconnect(immediate); + } + + private scheduleReconnect(immediate: boolean): void { + if (this.stopping || this.options === undefined || this.reconnectTimer !== undefined) return; + const base = immediate ? 0 : Math.min(1000 * 2 ** this.reconnectAttempts++, MAX_RECONNECT_DELAY_MS); + const delay = immediate ? 0 : Math.round(base * (0.75 + Math.random() * 0.5)); + this.reconnectTimer = setTimeout(() => { + this.reconnectTimer = undefined; + void this.connect().catch(() => { this.scheduleReconnect(false); }); + }, delay); + } + + private resetConnections(): void { + this.rootController?.abort(); + this.rootController = undefined; + for (const controller of this.requestControllers.values()) controller.abort(); + this.requestControllers.clear(); + for (const stream of this.streams.values()) { + stream.closeListener?.dispose(); + stream.bridge.dispose(); + } + this.streams.clear(); + this.assembler.clear(); + this.connectionResources.clear(); + this.management?.dispose(); + this.httpTunnel?.dispose(); + this.management = undefined; + this.httpTunnel = undefined; + this.httpTunnelOpening = false; + } + + private clearReconnectTimer(): void { + if (this.reconnectTimer === undefined) return; + clearTimeout(this.reconnectTimer); + this.reconnectTimer = undefined; + } + + private setState(state: RemoteControlState): void { + if (state === this._state) return; + this._state = state; + this._onDidChangeState.fire(state); + } +} + +function validateStartOptions(options: RemoteControlStartOptions): void { + const relay = new URL(options.relayBaseUrl); + if (relay.protocol !== 'http:' && relay.protocol !== 'https:') { + throw new Error('relay base URL must use HTTP or HTTPS'); + } + if (relay.username.length > 0 || relay.password.length > 0) { + throw new Error('relay base URL must not contain credentials'); + } + resolveLocalUrl(options.localBaseUrl, '/'); + if (options.alias.length === 0) throw new Error('remote control alias must not be empty'); + if ( + options.streamOpenTimeoutMs !== undefined && + (!Number.isFinite(options.streamOpenTimeoutMs) || options.streamOpenTimeoutMs <= 0) + ) { + throw new Error('stream open timeout must be a positive finite number'); + } +} + +function errorType(error: unknown): string { + return error instanceof Error ? error.name : typeof error; +} + +function streamErrorMessage(code: RemoteStreamErrorCode): string { + switch (code) { + case 'LOCAL_WS_FAILED': return 'failed to connect to local Kimi Code'; + case 'TUNNEL_STREAM_FAILED': return 'failed to connect relay stream'; + case 'TIMEOUT': return 'stream setup timed out'; + case 'UNKNOWN': return 'stream setup failed'; + } +} + +registerScopedService( + LifecycleScope.App, + IRemoteControlService, + RemoteControlService, + ScopeActivation.OnDemand, + 'remoteControl', +); diff --git a/packages/agent-core-v2/src/app/remoteControl/remoteControlTransport.ts b/packages/agent-core-v2/src/app/remoteControl/remoteControlTransport.ts new file mode 100644 index 0000000000..7b578cb0d0 --- /dev/null +++ b/packages/agent-core-v2/src/app/remoteControl/remoteControlTransport.ts @@ -0,0 +1,67 @@ +/** + * `remoteControl` domain — App-scope relay and localhost transport contract. + */ + +import type { IDisposable } from '#/_base/di/lifecycle'; +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; + +import type { LocalHttpRequest } from './protocol'; + +export interface RemoteSocket extends IDisposable { + send(data: string | Buffer, binary?: boolean): void; + close(code?: number, reason?: string): void; + ping(data?: Buffer): void; + pong(data?: Buffer): void; + onMessage(listener: (data: Buffer, binary: boolean) => void): IDisposable; + onClose(listener: (code: number, reason: string) => void): IDisposable; + onError(listener: (error: Error) => void): IDisposable; + onPing(listener: (data: Buffer) => void): IDisposable; + onPong(listener: (data: Buffer) => void): IDisposable; +} + +export interface LocalHttpResponse { + readonly statusCode: number; + readonly statusMessage: string; + readonly headers: Readonly>; + readonly streaming: boolean; + readonly body: AsyncIterable; +} + +export interface RemoteSocketBridge extends IDisposable { + close(code?: number, reason?: string): void; + onClose(listener: () => void): IDisposable; +} + +export interface IRemoteControlTransport { + readonly _serviceBrand: undefined; + connectManagement(relayBaseUrl: string, token: string, signal: AbortSignal): Promise; + connectHttpTunnel( + relayBaseUrl: string, + token: string, + deviceId: string, + signal: AbortSignal, + ): Promise; + connectTunnelStream( + relayBaseUrl: string, + streamId: string, + token: string, + signal: AbortSignal, + ): Promise; + connectLocalWebSocket( + localBaseUrl: string, + path: string, + headers: Readonly>, + localToken: string, + signal: AbortSignal, + ): Promise; + forwardLocalHttp( + localBaseUrl: string, + request: LocalHttpRequest, + localToken: string, + signal: AbortSignal, + ): Promise; + bridgeWebSockets(local: RemoteSocket, tunnel: RemoteSocket): RemoteSocketBridge; +} + +export const IRemoteControlTransport: ServiceIdentifier = + createDecorator('remoteControlTransport'); diff --git a/packages/agent-core-v2/src/app/remoteControl/remoteControlTransportService.ts b/packages/agent-core-v2/src/app/remoteControl/remoteControlTransportService.ts new file mode 100644 index 0000000000..5a867b018f --- /dev/null +++ b/packages/agent-core-v2/src/app/remoteControl/remoteControlTransportService.ts @@ -0,0 +1,300 @@ +/** + * `remoteControl` domain — Node HTTP and WebSocket transport implementation. + * + * Owns relay and localhost sockets at App scope and releases pending IO on disposal. + */ + +import http, { type IncomingHttpHeaders, type IncomingMessage } from 'node:http'; +import https from 'node:https'; + +import { WebSocket } from 'ws'; + +import { Disposable, combinedDisposable, toDisposable, type IDisposable } from '#/_base/di/lifecycle'; +import { LifecycleScope, ScopeActivation, registerScopedService } from '#/_base/di/scope'; + +import { resolveLocalUrl, sanitizeForwardHeaders, type LocalHttpRequest } from './protocol'; +import { + IRemoteControlTransport, + type LocalHttpResponse, + type RemoteSocket, + type RemoteSocketBridge, +} from './remoteControlTransport'; + +const BEARER_PROTOCOL_PREFIX = 'kimi-code.bearer.'; + +export class RemoteControlTransportService extends Disposable implements IRemoteControlTransport { + declare readonly _serviceBrand: undefined; + private readonly sockets = new Set(); + + connectManagement(base: string, token: string, signal: AbortSignal): Promise { + return this.connectRelay(base, '/v1/remote/create', token, signal); + } + + connectHttpTunnel(base: string, token: string, deviceId: string, signal: AbortSignal): Promise { + // The relay binds the tunnel to the registered device via the `device_id` + // query parameter; without it the socket is dropped right after upgrade. + return this.connectRelay( + base, + `/v1/remote/http?device_id=${encodeURIComponent(deviceId)}`, + token, + signal, + ); + } + + connectTunnelStream( + base: string, + streamId: string, + token: string, + signal: AbortSignal, + ): Promise { + return this.connectRelay(base, `/v1/remote/stream/${encodeURIComponent(streamId)}`, token, signal); + } + + connectLocalWebSocket( + base: string, + path: string, + headers: Readonly>, + localToken: string, + signal: AbortSignal, + ): Promise { + const url = resolveLocalUrl(base, path); + url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:'; + return this.connectSocket( + url, + [`${BEARER_PROTOCOL_PREFIX}${localToken}`], + { + ...sanitizeForwardHeaders(headers), + Authorization: `Bearer ${localToken}`, + }, + signal, + ); + } + + forwardLocalHttp( + base: string, + request: LocalHttpRequest, + localToken: string, + signal: AbortSignal, + ): Promise { + const url = resolveLocalUrl(base, request.path); + const requestFn = url.protocol === 'https:' ? https.request : http.request; + return new Promise((resolve, reject) => { + const req = requestFn( + url, + { + method: request.method, + headers: { + ...sanitizeForwardHeaders(request.headers), + authorization: `Bearer ${localToken}`, + 'content-length': String(request.body.byteLength), + }, + signal, + }, + (response) => { resolve(toLocalHttpResponse(response)); }, + ); + req.once('error', reject); + if (request.body.byteLength > 0) req.write(request.body); + req.end(); + }); + } + + bridgeWebSockets(local: RemoteSocket, tunnel: RemoteSocket): RemoteSocketBridge { + let closed = false; + let subscriptions: IDisposable | undefined; + const listeners = new Set<() => void>(); + const closeBoth = (code = 1000, reason = ''): void => { + if (closed) return; + closed = true; + local.close(code, reason); + tunnel.close(code, reason); + subscriptions?.dispose(); + for (const listener of listeners) listener(); + listeners.clear(); + }; + subscriptions = combinedDisposable( + local.onMessage((data, binary) => { tunnel.send(data, binary); }), + tunnel.onMessage((data, binary) => { local.send(data, binary); }), + local.onPing((data) => { tunnel.ping(data); }), + tunnel.onPing((data) => { local.ping(data); }), + local.onPong((data) => { tunnel.pong(data); }), + tunnel.onPong((data) => { local.pong(data); }), + local.onClose(closeBoth), + tunnel.onClose(closeBoth), + ); + return { + close: closeBoth, + onClose: (listener) => { + if (closed) { + listener(); + return toDisposable(() => {}); + } + listeners.add(listener); + return toDisposable(() => { listeners.delete(listener); }); + }, + dispose: closeBoth, + }; + } + + override dispose(): void { + for (const socket of this.sockets) socket.dispose(); + this.sockets.clear(); + super.dispose(); + } + + private connectRelay( + base: string, + path: string, + token: string, + signal: AbortSignal, + ): Promise { + // `path` is origin-absolute (`/v1/remote/...`), so `new URL(path, base)` + // would discard any mount prefix in the base URL (e.g. `/coding-relay`). + // Splice the relay paths under the base pathname instead. + const rel = new URL(path, 'http://relay.invalid'); + const url = new URL(base); + url.pathname = `${url.pathname.replace(/\/+$/, '')}${rel.pathname}`; + url.search = rel.search; + url.protocol = url.protocol === 'http:' ? 'ws:' : 'wss:'; + // The relay authenticates via the standard Authorization header; the bearer + // subprotocol convention is only for the localhost hop, where the local + // server echoes the protocol and the `ws` client requires that echo. + return this.connectSocket(url, [], { Authorization: `Bearer ${token}` }, signal); + } + + private async connectSocket( + url: URL, + protocols: string[], + headers: Readonly>, + signal: AbortSignal, + ): Promise { + const ws = new WebSocket(url, protocols, { headers, autoPong: false }); + const socket = new NodeRemoteSocket(ws, () => this.sockets.delete(socket)); + this.sockets.add(socket); + try { + await waitForOpen(ws, signal); + return socket; + } catch (error) { + socket.dispose(); + throw error; + } + } +} + +class NodeRemoteSocket implements RemoteSocket { + private disposed = false; + + constructor( + private readonly ws: WebSocket, + private readonly onDispose: () => void, + ) {} + + send(data: string | Buffer, binary = Buffer.isBuffer(data)): void { + if (this.ws.readyState === WebSocket.OPEN) this.ws.send(data, { binary }); + } + + close(code = 1000, reason = ''): void { + if (this.ws.readyState === WebSocket.OPEN || this.ws.readyState === WebSocket.CONNECTING) { + this.ws.close(code, reason); + } + } + + ping(data?: Buffer): void { + if (this.ws.readyState === WebSocket.OPEN) this.ws.ping(data); + } + + pong(data?: Buffer): void { + if (this.ws.readyState === WebSocket.OPEN) this.ws.pong(data); + } + + onMessage(listener: (data: Buffer, binary: boolean) => void): IDisposable { + const handler = (data: WebSocket.RawData, binary: boolean): void => { + listener(toBuffer(data), binary); + }; + this.ws.on('message', handler); + return toDisposable(() => this.ws.off('message', handler)); + } + + onClose(listener: (code: number, reason: string) => void): IDisposable { + const handler = (code: number, reason: Buffer): void => { + listener(code, reason.toString()); + }; + this.ws.on('close', handler); + return toDisposable(() => this.ws.off('close', handler)); + } + + onError(listener: (error: Error) => void): IDisposable { + this.ws.on('error', listener); + return toDisposable(() => this.ws.off('error', listener)); + } + + onPing(listener: (data: Buffer) => void): IDisposable { + this.ws.on('ping', listener); + return toDisposable(() => this.ws.off('ping', listener)); + } + + onPong(listener: (data: Buffer) => void): IDisposable { + this.ws.on('pong', listener); + return toDisposable(() => this.ws.off('pong', listener)); + } + + dispose(): void { + if (this.disposed) return; + this.disposed = true; + this.onDispose(); + this.ws.terminate(); + } +} + +function waitForOpen(ws: WebSocket, signal: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const cleanup = (): void => { + ws.off('open', onOpen); + ws.off('error', onError); + signal.removeEventListener('abort', onAbort); + }; + const onOpen = (): void => { cleanup(); resolve(); }; + const onError = (error: Error): void => { cleanup(); reject(error); }; + const onAbort = (): void => { cleanup(); reject(signal.reason ?? new Error('aborted')); }; + if (signal.aborted) { + onAbort(); + return; + } + ws.once('open', onOpen); + ws.once('error', onError); + signal.addEventListener('abort', onAbort, { once: true }); + }); +} + +function toLocalHttpResponse(response: IncomingMessage): LocalHttpResponse { + const headers = normalizeHeaders(response.headers); + const contentType = response.headers['content-type'] ?? ''; + const streaming = contentType.toLowerCase().includes('text/event-stream') || + response.headers['transfer-encoding'] !== undefined; + return { + statusCode: response.statusCode ?? 502, + statusMessage: response.statusMessage ?? 'Bad Gateway', + headers, + streaming, + body: response as AsyncIterable, + }; +} + +function normalizeHeaders( + headers: IncomingHttpHeaders, +): Record { + return { ...headers }; +} + +function toBuffer(data: WebSocket.RawData): Buffer { + if (Array.isArray(data)) return Buffer.concat(data); + if (data instanceof ArrayBuffer) return Buffer.from(data); + return Buffer.from(data); +} + +registerScopedService( + LifecycleScope.App, + IRemoteControlTransport, + RemoteControlTransportService, + ScopeActivation.OnDemand, + 'remoteControl', +); diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 650e5e0beb..91bf400fa0 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -375,6 +375,13 @@ export * from '#/agent/tools/ask-user-question/ask-user-question'; import '#/agent/tools/ask-user-question/askUserQuestionTool'; export * from '#/app/gateway/gateway'; export * from '#/app/gateway/gatewayService'; +import '#/app/remoteControl/flag'; +export * from '#/app/remoteControl/flag'; +export * from '#/app/remoteControl/protocol'; +export * from '#/app/remoteControl/remoteControl'; +export * from '#/app/remoteControl/remoteControlTransport'; +import '#/app/remoteControl/remoteControlService'; +import '#/app/remoteControl/remoteControlTransportService'; export * from '#/session/workspaceContext/workspaceContext'; export * from '#/session/workspaceContext/workspaceContextService'; diff --git a/packages/agent-core-v2/test/app/remoteControl/remoteControl.test.ts b/packages/agent-core-v2/test/app/remoteControl/remoteControl.test.ts new file mode 100644 index 0000000000..e01b914f9b --- /dev/null +++ b/packages/agent-core-v2/test/app/remoteControl/remoteControl.test.ts @@ -0,0 +1,499 @@ +/** + * Remote Control boundary scenarios — validates protocol codecs, localhost URL/header isolation, + * device tunnel lifecycle, stream bridging, and reconnect behavior through DI-resolved services. + * OAuth and network transports are stubbed; protocol helpers and the Node bridge are real. + * Run with: pnpm --filter @moonshot-ai/agent-core-v2 exec vitest run test/app/remoteControl/remoteControl.test.ts + */ + +import { mkdtemp, rm } from 'node:fs/promises'; +import { createServer } from 'node:http'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { WebSocketServer } from 'ws'; + +import { DisposableStore, toDisposable, type IDisposable } from '#/_base/di/lifecycle'; +import { createServices, type TestInstantiationService } from '#/_base/di/test'; +import { ILogService } from '#/_base/log/log'; +import { IOAuthService } from '#/app/auth/auth'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { IFlagService } from '#/app/flag/flag'; +import { + BAD_GATEWAY_RESPONSE, + HttpRequestAssembler, + parseRawHttpRequest, + resolveLocalUrl, + sanitizeForwardHeaders, + type HttpTunnelMessage, + type LocalHttpRequest, +} from '#/app/remoteControl/protocol'; +import { IRemoteControlService } from '#/app/remoteControl/remoteControl'; +import { RemoteControlService } from '#/app/remoteControl/remoteControlService'; +import { + IRemoteControlTransport, + type LocalHttpResponse, + type RemoteSocket, + type RemoteSocketBridge, +} from '#/app/remoteControl/remoteControlTransport'; +import { RemoteControlTransportService } from '#/app/remoteControl/remoteControlTransportService'; +import { stubLog } from '../../_base/log/stubs'; + +class FakeSocket implements RemoteSocket { + readonly sent: Array = []; + readonly closeCalls: Array<{ code: number; reason: string }> = []; + closed = false; + private readonly messages = new Set<(data: Buffer, binary: boolean) => void>(); + private readonly closes = new Set<(code: number, reason: string) => void>(); + private readonly errors = new Set<(error: Error) => void>(); + private readonly pings = new Set<(data: Buffer) => void>(); + private readonly pongs = new Set<(data: Buffer) => void>(); + + send(data: string | Buffer): void { this.sent.push(data); } + close(code = 1000, reason = ''): void { + if (this.closed) return; + this.closed = true; + this.closeCalls.push({ code, reason }); + for (const listener of this.closes) listener(code, reason); + } + ping(data = Buffer.alloc(0)): void { for (const listener of this.pings) listener(data); } + pong(data = Buffer.alloc(0)): void { for (const listener of this.pongs) listener(data); } + onMessage(listener: (data: Buffer, binary: boolean) => void): IDisposable { + this.messages.add(listener); return toDisposable(() => this.messages.delete(listener)); + } + onClose(listener: (code: number, reason: string) => void): IDisposable { + this.closes.add(listener); return toDisposable(() => this.closes.delete(listener)); + } + onError(listener: (error: Error) => void): IDisposable { + this.errors.add(listener); return toDisposable(() => this.errors.delete(listener)); + } + onPing(listener: (data: Buffer) => void): IDisposable { + this.pings.add(listener); return toDisposable(() => this.pings.delete(listener)); + } + onPong(listener: (data: Buffer) => void): IDisposable { + this.pongs.add(listener); return toDisposable(() => this.pongs.delete(listener)); + } + emit(value: unknown): void { + const data = Buffer.from(JSON.stringify(value)); + for (const listener of this.messages) listener(data, false); + } + emitClose(code = 1000, reason = ''): void { + this.closed = true; + for (const listener of this.closes) listener(code, reason); + } + dispose(): void { this.closed = true; } +} + +class FakeBridge implements RemoteSocketBridge { + readonly closeCalls: Array<{ code: number; reason: string }> = []; + private readonly listeners = new Set<() => void>(); + + close(code = 1000, reason = ''): void { + this.closeCalls.push({ code, reason }); + this.emitClose(); + } + onClose(listener: () => void): IDisposable { + this.listeners.add(listener); + return toDisposable(() => { this.listeners.delete(listener); }); + } + emitClose(): void { + for (const listener of this.listeners) listener(); + this.listeners.clear(); + } + dispose(): void { this.close(); } +} + +class FakeTransport implements IRemoteControlTransport { + readonly _serviceBrand = undefined; + readonly management = new FakeSocket(); + readonly http = new FakeSocket(); + readonly order: string[] = []; + readonly bridges: FakeBridge[] = []; + local = new FakeSocket(); + tunnel = new FakeSocket(); + response: LocalHttpResponse = responseOf(200, [Buffer.from('ok')]); + managementError = false; + forwardError = false; + localError = false; + tunnelError = false; + bridgeError = false; + forwarded: Array<{ request: LocalHttpRequest; token: string }> = []; + + connectManagement(): Promise { + this.order.push('management'); + return this.managementError + ? Promise.reject(new Error('relay unavailable')) + : Promise.resolve(this.management); + } + connectHttpTunnel(): Promise { this.order.push('http'); return Promise.resolve(this.http); } + connectTunnelStream(): Promise { + this.order.push('tunnel'); + return this.tunnelError ? Promise.reject(new Error('tunnel failed')) : Promise.resolve(this.tunnel); + } + connectLocalWebSocket(): Promise { + this.order.push('local'); + return this.localError ? Promise.reject(new Error('local failed')) : Promise.resolve(this.local); + } + forwardLocalHttp(_base: string, request: LocalHttpRequest, token: string): Promise { + this.forwarded.push({ request, token }); + return this.forwardError ? Promise.reject(new Error('local unavailable')) : Promise.resolve(this.response); + } + bridgeWebSockets(): RemoteSocketBridge { + this.order.push('bridge'); + if (this.bridgeError) throw new Error('bridge failed'); + const bridge = new FakeBridge(); + this.bridges.push(bridge); + return bridge; + } +} + +function responseOf(statusCode: number, chunks: Buffer[], streaming = false): LocalHttpResponse { + return { + statusCode, + statusMessage: statusCode === 200 ? 'OK' : 'Error', + headers: { 'content-type': streaming ? 'text/event-stream' : 'application/json' }, + streaming, + body: (async function* () { for (const chunk of chunks) yield chunk; })(), + }; +} + +function sentJson(socket: FakeSocket): unknown[] { + return socket.sent.map((value) => JSON.parse(value.toString())); +} + +async function flushMicrotasks(): Promise { + for (let i = 0; i < 8; i++) await Promise.resolve(); +} + +describe('remote control protocol', () => { + it('assembles concurrent request fragments and strips remote credentials', () => { + const assembler = new HttpRequestAssembler(); + const part = (requestId: string, raw: string, isLast: boolean): HttpTunnelMessage => ({ + request_id: requestId, + type: 'request', + is_last: isLast, + body_base64: Buffer.from(raw).toString('base64'), + }); + expect(assembler.push(part('a', 'POST /a HTTP/1.1\r\nAuthorization: Bearer remote\r\nContent-Length: 3\r\n\r\n', false))).toBeUndefined(); + const requestB = assembler.push(part('b', 'GET /b?q=1 HTTP/1.1\r\nConnection: close\r\n\r\n', true)); + const requestA = assembler.push(part('a', 'abc', true)); + expect(requestB).toMatchObject({ method: 'GET', path: '/b?q=1', headers: {} }); + expect(requestA).toMatchObject({ method: 'POST', path: '/a', headers: {}, body: Buffer.from('abc') }); + }); + + it('rejects malformed raw HTTP requests', () => { + expect(() => parseRawHttpRequest(Buffer.from('not HTTP'))).toThrow(); + }); + + it('uses the protocol-defined empty 502 response', () => { + expect(BAD_GATEWAY_RESPONSE.toString()).toBe( + 'HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n', + ); + }); +}); + +describe('remote control localhost boundary', () => { + it('resolves a rooted path within the local origin', () => { + expect(resolveLocalUrl('http://127.0.0.1:4321', '/api/v1/sessions?q=1').href).toBe( + 'http://127.0.0.1:4321/api/v1/sessions?q=1', + ); + }); + + it.each(['https://example.test/steal', '//example.test/steal', 'api/v1/sessions'])( + 'rejects non-local path %s before forwarding', + (path) => { + expect(() => resolveLocalUrl('http://127.0.0.1:4321', path)).toThrow(); + }, + ); + + it('strips Origin and WebSocket handshake headers before localhost forwarding', () => { + expect(sanitizeForwardHeaders({ + Origin: 'https://remote.example.test', + Authorization: 'Bearer remote', + 'Sec-WebSocket-Key': 'remote-key', + 'Sec-WebSocket-Extensions': 'permessage-deflate', + 'Sec-WebSocket-Protocol': 'remote-protocol', + 'Sec-WebSocket-Version': '13', + 'User-Agent': 'browser', + })).toEqual({ 'User-Agent': 'browser' }); + }); + + it('rejects an absolute WebSocket URL before opening a socket', () => { + const transport = new RemoteControlTransportService(); + expect(() => transport.connectLocalWebSocket( + 'http://127.0.0.1:4321', + 'https://example.test/steal', + {}, + 'local-token', + new AbortController().signal, + )).toThrow(); + transport.dispose(); + }); + + it('rejects an absolute HTTP URL before attaching the local token', () => { + const transport = new RemoteControlTransportService(); + expect(() => transport.forwardLocalHttp( + 'http://127.0.0.1:4321', + { method: 'GET', path: 'https://example.test/steal', headers: {}, body: Buffer.alloc(0) }, + 'local-token', + new AbortController().signal, + )).toThrow(); + transport.dispose(); + }); + + it('closes both bridge sockets with the requested code and reason', () => { + const transport = new RemoteControlTransportService(); + const local = new FakeSocket(); + const tunnel = new FakeSocket(); + const bridge = transport.bridgeWebSockets(local, tunnel); + bridge.close(4001, 'browser_closed'); + expect(local.closeCalls).toEqual([{ code: 4001, reason: 'browser_closed' }]); + expect(tunnel.closeCalls).toEqual([{ code: 4001, reason: 'browser_closed' }]); + transport.dispose(); + }); +}); + +describe('remote control relay transport', () => { + it('connects relay sockets under the base URL mount prefix', async () => { + const upgrades: Array<{ url: string | undefined; authorization: string | undefined }> = []; + const wss = new WebSocketServer({ noServer: true }); + const httpServer = createServer(); + httpServer.on('upgrade', (req, socket, head) => { + upgrades.push({ url: req.url, authorization: req.headers.authorization }); + wss.handleUpgrade(req, socket, head, (ws) => { wss.emit('connection', ws, req); }); + }); + await new Promise((resolve) => { httpServer.listen(0, '127.0.0.1', resolve); }); + try { + const address = httpServer.address(); + if (address === null || typeof address === 'string') throw new Error('listener has no port'); + const transport = new RemoteControlTransportService(); + const socket = await transport.connectManagement( + `http://127.0.0.1:${String(address.port)}/coding-relay`, + 'relay-token', + new AbortController().signal, + ); + expect(upgrades).toEqual([{ + url: '/coding-relay/v1/remote/create', + authorization: 'Bearer relay-token', + }]); + socket.dispose(); + transport.dispose(); + } finally { + wss.close(); + await new Promise((resolve) => { httpServer.close(() => { resolve(); }); }); + } + }); +}); + +describe('Remote Control service lifecycle', () => { + let disposables: DisposableStore; + let ix: TestInstantiationService; + let transport: FakeTransport; + let homeDir: string; + let service: IRemoteControlService; + const localToken = 'local-secret'; + const deviceToken = 'device-secret'; + + beforeEach(async () => { + disposables = new DisposableStore(); + transport = new FakeTransport(); + homeDir = await mkdtemp(join(tmpdir(), 'remote-control-test-')); + ix = createServices(disposables, { + additionalServices: (reg) => { + reg.defineInstance(IRemoteControlTransport, transport); + reg.definePartialInstance(IOAuthService, { + getCachedAccessToken: () => Promise.resolve(deviceToken), + }); + reg.definePartialInstance(IBootstrapService, { + homeDir, + platform: 'darwin', + clientIdentity: { productName: 'test-host', version: '1.2.3', platform: 'test_platform' }, + }); + reg.definePartialInstance(IFlagService, { enabled: () => true }); + reg.defineInstance(ILogService, stubLog()); + reg.define(IRemoteControlService, RemoteControlService); + }, + strict: true, + }); + service = ix.get(IRemoteControlService); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + disposables.dispose(); + await rm(homeDir, { recursive: true, force: true }); + }); + + function options() { + return { + relayBaseUrl: 'https://relay.example.test', + localBaseUrl: 'http://127.0.0.1:4321', + alias: 'test-device', + getLocalToken: () => localToken, + }; + } + + async function startAndAck(): Promise { + await service.start(options()); + transport.management.emit({ type: 'register_ack', payload: { success: true } }); + await flushMicrotasks(); + } + + it('sends device metadata without embedding the OAuth token when registration starts', async () => { + await service.start(options()); + const register = sentJson(transport.management)[0] as { payload: Record }; + expect(register.payload).toMatchObject({ + alias: 'test-device', + platform: 'darwin', + client_version: '1.2.3', + local_base_url: 'http://127.0.0.1:4321', + }); + expect(JSON.stringify(register)).not.toContain(deviceToken); + }); + + it('enters online after registration opens the HTTP tunnel', async () => { + const states: string[] = []; + service.onDidChangeState((state) => states.push(state)); + await startAndAck(); + expect(service.state).toBe('online'); + expect(states).toEqual(['connecting', 'online']); + }); + + it('opens only one HTTP tunnel when register_ack is repeated', async () => { + await service.start(options()); + transport.management.emit({ type: 'register_ack', payload: { success: true } }); + transport.management.emit({ type: 'register_ack', payload: { success: true } }); + await flushMicrotasks(); + expect(transport.order.filter((step) => step === 'http')).toHaveLength(1); + }); + + it('forwards a reassembled HTTP upload with only the local credential', async () => { + await startAndAck(); + const raw = 'POST /api/test HTTP/1.1\r\nAuthorization: Bearer remote\r\n\r\nbody'; + transport.http.emit({ + request_id: 'r1', type: 'request', is_last: true, + body_base64: Buffer.from(raw).toString('base64'), + }); + await flushMicrotasks(); + expect(transport.forwarded[0]).toMatchObject({ token: localToken }); + expect(transport.forwarded[0]!.request.headers).toEqual({}); + }); + + it('emits response chunks when localhost returns a streaming response', async () => { + transport.response = responseOf(200, [Buffer.from('one'), Buffer.from('two')], true); + await startAndAck(); + transport.http.emit({ + request_id: 'r1', type: 'request', is_last: true, + body_base64: Buffer.from('GET /events HTTP/1.1\r\n\r\n').toString('base64'), + }); + await flushMicrotasks(); + const messages = sentJson(transport.http) as HttpTunnelMessage[]; + expect(messages.map((message) => [message.type, message.is_last])).toEqual([ + ['response_chunk', false], + ['response_chunk', false], + ['response_chunk', false], + ['response_chunk', true], + ]); + }); + + it('returns a synthetic 502 when localhost forwarding fails', async () => { + transport.forwardError = true; + await startAndAck(); + transport.http.emit({ + request_id: 'r2', type: 'request', is_last: true, + body_base64: Buffer.from('GET / HTTP/1.1\r\n\r\n').toString('base64'), + }); + await flushMicrotasks(); + const response = sentJson(transport.http)[0] as HttpTunnelMessage; + expect(Buffer.from(response.body_base64, 'base64')).toEqual(BAD_GATEWAY_RESPONSE); + }); + + it('opens the relay stream before localhost when handling open_ws', async () => { + await startAndAck(); + transport.management.emit({ + type: 'open_ws', + payload: { stream_id: 's1', path: '/api/v1/ws', headers: { Authorization: 'remote' } }, + }); + await flushMicrotasks(); + expect(transport.order.slice(-3)).toEqual(['tunnel', 'local', 'bridge']); + }); + + it.each([ + ['localError', 'LOCAL_WS_FAILED'], + ['tunnelError', 'TUNNEL_STREAM_FAILED'], + ['bridgeError', 'UNKNOWN'], + ] as const)('reports %s as %s when stream setup fails', async (failure, errorCode) => { + transport[failure] = true; + await startAndAck(); + transport.management.emit({ + type: 'open_ws', payload: { stream_id: 'failed', path: '/api/v1/ws', headers: {} }, + }); + await flushMicrotasks(); + expect(sentJson(transport.management).at(-1)).toMatchObject({ + type: 'open_ws_result', + payload: { stream_id: 'failed', success: false, error_code: errorCode }, + }); + }); + + it('passes relay close code and reason to the active stream bridge', async () => { + await startAndAck(); + transport.management.emit({ + type: 'open_ws', payload: { stream_id: 's1', path: '/api/v1/ws', headers: {} }, + }); + await flushMicrotasks(); + transport.management.emit({ + type: 'close_ws', + payload: { stream_id: 's1', close_code: 4001, reason: 'browser_closed' }, + }); + await flushMicrotasks(); + expect(transport.bridges[0]!.closeCalls).toEqual([{ code: 4001, reason: 'browser_closed' }]); + }); + + it('forgets a stream after the bridge closes naturally', async () => { + await startAndAck(); + transport.management.emit({ + type: 'open_ws', payload: { stream_id: 's1', path: '/api/v1/ws', headers: {} }, + }); + await flushMicrotasks(); + transport.bridges[0]!.emitClose(); + transport.management.emit({ + type: 'close_ws', payload: { stream_id: 's1', reason: 'browser_closed' }, + }); + await flushMicrotasks(); + expect(transport.bridges[0]!.closeCalls).toEqual([]); + }); + + it('resolves offline when the first relay connection is temporarily unavailable', async () => { + transport.managementError = true; + await expect(service.start(options())).resolves.toBeUndefined(); + expect(service.state).toBe('offline'); + }); + + it('rejects an invalid relay URL before starting the tunnel', async () => { + await expect(service.start({ ...options(), relayBaseUrl: 'not a URL' })).rejects.toThrow(); + expect(transport.order).toEqual([]); + }); + + it('sends local_server_stopped when the service stops', async () => { + await startAndAck(); + await service.stop('local_server_stopped'); + expect(sentJson(transport.management).at(-1)).toEqual({ + type: 'disconnect', payload: { reason: 'local_server_stopped' }, + }); + }); + + it('closes the management channel when an invalid message arrives', async () => { + await service.start(options()); + transport.management.emit({ type: 'unknown' }); + expect(transport.management.closed).toBe(true); + }); + + it('does not log credentials when an invalid message arrives', async () => { + const log = ix.get(ILogService); + const warn = vi.spyOn(log, 'warn'); + await service.start(options()); + transport.management.emit({ type: 'unknown', token: deviceToken }); + expect(JSON.stringify(warn.mock.calls)).not.toContain(deviceToken); + expect(JSON.stringify(warn.mock.calls)).not.toContain(localToken); + }); +}); diff --git a/packages/kap-server/src/start.ts b/packages/kap-server/src/start.ts index 7e92e93ca7..70811e2b05 100644 --- a/packages/kap-server/src/start.ts +++ b/packages/kap-server/src/start.ts @@ -11,6 +11,7 @@ import { bootstrap, IConfigService, IProviderDiscoveryService, + IRemoteControlService, IWorkspaceService, logSeed, resolveConfigPath, @@ -38,6 +39,7 @@ import { type ServerLogger, type ServerLogLevel, } from './services/pinoLoggerService'; +import { hostname } from 'node:os'; import { join } from 'node:path'; import type { Socket } from 'node:net'; import type { IncomingMessage } from 'node:http'; @@ -122,6 +124,11 @@ export interface ServerStartOptions { * unset unless a second, distinct RPC credential is genuinely needed. */ readonly rpcToken?: string; + /** Optional relay endpoint for the experimental Remote Control tunnel. */ + readonly remoteControl?: { + readonly relayBaseUrl: string; + readonly alias?: string; + }; /** Extra scope seeds applied at bootstrap (e.g. a host-provided `ISessionModelResolver`). */ readonly seeds?: ScopeSeed; /** @@ -349,7 +356,9 @@ export async function startServer(opts: ServerStartOptions): Promise => { + await remoteControlService?.stop('local_server_stopped'); await app.close(); authFailureLimiter?.dispose(); modelCatalogRefreshScheduler.dispose(); @@ -610,6 +619,25 @@ export async function startServer(opts: ServerStartOptions): Promise authTokenService.getToken(), + }); + } catch (error) { + try { + await close(); + } catch { + // best-effort cleanup; the Remote Control start error is what matters + } + throw error; + } + } + void modelCatalogRefreshScheduler.start().catch((error) => { logger.warn( { err: error instanceof Error ? error.message : String(error) }, @@ -620,6 +648,12 @@ export async function startServer(opts: ServerStartOptions): Promise { expect(() => core.accessor.get(IBootstrapService)).toThrow(); expect(await listLiveServerInstances(home)).toEqual([]); }); + + it('starts Remote Control after listen with the bound port and stops it before Fastify', async () => { + home = await mkdtemp(join(tmpdir(), 'kimi-server-v2-remote-control-')); + const events: string[] = []; + let startOptions: RemoteControlStartOptions | undefined; + const remoteControl: IRemoteControlService = { + _serviceBrand: undefined, + state: 'offline', + onDidChangeState: () => ({ dispose: () => {} }), + start: (options) => { + startOptions = options; + events.push('remote-start'); + return Promise.resolve(); + }, + stop: (reason) => { + events.push(`remote-stop:${reason}`); + return Promise.resolve(); + }, + }; + const authTokenService: IAuthTokenService = { + _serviceBrand: undefined, + getToken: () => 'local-token', + isValid: (candidate) => Promise.resolve(candidate === 'local-token'), + }; + server = await startServer({ + hostIdentity: TEST_HOST_IDENTITY, + host: '0.0.0.0', + port: 0, + homeDir: home, + insecureNoTls: true, + logLevel: 'silent', + authTokenService, + remoteControl: { relayBaseUrl: 'https://relay.example.test', alias: 'test-device' }, + seeds: [[IRemoteControlService as ServiceIdentifier, remoteControl]], + }); + const closeFastify = server.app.close.bind(server.app); + const appClose = server.app as unknown as { close(): Promise }; + appClose.close = async () => { + events.push('fastify-close'); + await closeFastify(); + }; + + expect(startOptions).toMatchObject({ + relayBaseUrl: 'https://relay.example.test', + alias: 'test-device', + localBaseUrl: `http://127.0.0.1:${String(server.port)}`, + }); + expect(startOptions?.getLocalToken()).toBe('local-token'); + expect(events).toEqual(['remote-start']); + + await server.close(); + server = undefined; + expect(events).toEqual([ + 'remote-start', + 'remote-stop:local_server_stopped', + 'fastify-close', + ]); + }); }); function silentLogger() { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 72d8deb208..9f9823f499 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -758,6 +758,9 @@ importers: undici: specifier: ^7.27.1 version: 7.27.1 + ws: + specifier: ^8.18.0 + version: 8.20.0 yauzl: specifier: ^3.3.0 version: 3.3.0 @@ -792,6 +795,9 @@ importers: '@types/tar': specifier: ^7.0.87 version: 7.0.87 + '@types/ws': + specifier: ^8.18.0 + version: 8.18.1 '@types/yauzl': specifier: ^2.10.3 version: 2.10.3