basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6/**7 * Client-hosted MCP over the daemon WS (issue #5626, Phase 2 "reverse tool8 * channel").9 *10 * A connected WS client (the Chrome extension) hosts an MCP server (its11 * browser tools) that the daemon's agent calls. The agent's MCP client side is12 * the existing `SdkControlClientTransport`; this module is the daemon-WS glue13 * that carries the same `mcp_message` JSON-RPC frames over the WS rather than14 * the SDK subprocess control plane.15 *16 * Per-connection lifecycle:17 *   - `mcp_register { server }`     → register an SDK-type runtime MCP server18 *                                     whose `sendSdkMcpMessage` pushes19 *                                     `mcp_message` frames down THIS WS.20 *   - `mcp_message { id, server, payload }` (client→daemon) → resolve the21 *                                     correlated pending request.22 *   - `mcp_unregister { server }`   → remove the runtime server + reject23 *                                     pending.24 *   - WS close                      → tear down all of the connection's25 *                                     servers + reject pending.26 *27 * Wiring status (see the architecture note in `05-daemon-direct-architecture.md`28 * and the PR notes): the WS framing + correlation is fully wired through the29 * real `ClientMcpRegistrar` / `SdkControlClientTransport` round-trip. The deep30 * hookup into the agent's live `McpClientManager` is injected via a31 * {@link ClientMcpServerProvider}. In the current daemon the `McpClientManager`32 * lives in the ACP child process while this WS lives in the parent, so the33 * provider is wired only when one is supplied (the round-trip test supplies a34 * real in-process manager). When absent, registration is rejected with a35 * structured `not_wired` error so the contract stays honest.36 */37import { ClientMcpRegistrar, } from '@qwen-code/qwen-code-core';38import { isValidServerName } from '../validate-server-name.js';39/** WS frame discriminators owned by this module. */40export const CLIENT_MCP_FRAME_TYPES = {41    register: 'mcp_register',42    message: 'mcp_message',43    unregister: 'mcp_unregister',44};45/**46 * Upper bound on client-hosted MCP servers per WS connection. Caps the runtime47 * MCP add/discovery a single client can drive so a misbehaving (or hostile)48 * client can't register an unbounded number of servers on one connection.49 */50const MAX_SERVERS_PER_CONNECTION = 10;51/**52 * Per-WS-connection holder for client-hosted MCP servers. One instance per53 * connection; disposed on WS close.54 */55export class ClientMcpWsConnection {56    sendFrame;57    provider;58    registrar;59    disposed = false;60    constructor(sendFrame, provider) {61        this.sendFrame = sendFrame;62        this.provider = provider;63        this.registrar = new ClientMcpRegistrar({64            sendFrame: (frame) => {65                this.sendFrame({66                    type: CLIENT_MCP_FRAME_TYPES.message,67                    id: frame.id,68                    server: frame.server,69                    payload: frame.payload,70                });71            },72        });73    }74    /**75     * Route a parsed inbound frame. Returns a structured result the WS layer can76     * turn into an ack/error reply (or ignore). Never throws — protocol errors77     * are returned as `{ kind: 'error' }`.78     */79    async handleFrame(frame) {80        if (this.disposed) {81            return { kind: 'error', code: 'closed', message: 'connection closed' };82        }83        switch (frame.type) {84            case CLIENT_MCP_FRAME_TYPES.register:85                return this.handleRegister(frame.server);86            case CLIENT_MCP_FRAME_TYPES.unregister:87                return this.handleUnregister(frame.server);88            case CLIENT_MCP_FRAME_TYPES.message:89                return this.handleMessage(frame.id, frame.payload);90            default:91                return {92                    kind: 'ignored',93                    reason: `unknown client-mcp frame type: ${String(frame.type)}`,94                };95        }96    }97    /** Whether a frame's `type` is one this module owns. */98    static isClientMcpFrameType(type) {99        return (type === CLIENT_MCP_FRAME_TYPES.register ||100            type === CLIENT_MCP_FRAME_TYPES.message ||101            type === CLIENT_MCP_FRAME_TYPES.unregister);102    }103    async handleRegister(server) {104        if (!isValidServerName(server)) {105            return {106                kind: 'error',107                code: 'invalid_server_name',108                message: 'server must be ≤256 chars, alphanumeric + underscore/hyphen, and not a reserved JS property name',109            };110        }111        if (this.registrar.hasServer(server)) {112            return {113                kind: 'error',114                code: 'already_registered',115                message: `server '${server}' is already registered on this connection`,116            };117        }118        // Cap the number of servers a single connection can register so a client119        // can't drive unbounded runtime-MCP add/discovery (DoS guard).120        if (this.registrar.serverCount() >= MAX_SERVERS_PER_CONNECTION) {121            return {122                kind: 'error',123                code: 'too_many_servers',124                message: `connection has reached the maximum of ${MAX_SERVERS_PER_CONNECTION} registered MCP servers`,125            };126        }127        if (!this.provider) {128            return {129                kind: 'error',130                code: 'not_wired',131                message: 'client_mcp_over_ws is advertised but no McpClientManager provider is wired into this daemon process',132            };133        }134        // Advertise to the registrar BEFORE registering so the SDK discovery135        // handshake (which the provider triggers synchronously) can route frames.136        this.registrar.registerServer(server);137        try {138            const { toolCount } = await this.provider.registerClientMcpServer(server, this.registrar.sendSdkMcpMessage);139            // The WS may have closed (dispose() ran) while we awaited the provider140            // round-trip. dispose() snapshots its server set before this register141            // resolves, so the provider would otherwise be left holding a zombie142            // runtime MCP server. Re-check and tear it back down.143            if (this.disposed) {144                this.registrar.unregisterServer(server);145                await this.provider.unregisterClientMcpServer(server);146                return {147                    kind: 'error',148                    code: 'closed',149                    message: 'connection disposed during register',150                };151            }152            return { kind: 'registered', server, toolCount };153        }154        catch (err) {155            // Roll back the registrar advertisement on failure.156            this.registrar.unregisterServer(server);157            return {158                kind: 'error',159                code: 'register_failed',160                message: err instanceof Error ? err.message : String(err),161            };162        }163    }164    async handleUnregister(server) {165        if (this.disposed) {166            return { kind: 'unregistered', server: String(server) };167        }168        if (!isValidServerName(server)) {169            return {170                kind: 'error',171                code: 'invalid_server_name',172                message: 'server name is invalid',173            };174        }175        const existed = this.registrar.unregisterServer(server);176        if (existed && this.provider) {177            try {178                await this.provider.unregisterClientMcpServer(server);179            }180            catch {181                // Best-effort teardown — the registrar already rejected pending.182            }183        }184        // Idempotent on purpose: report `unregistered` whether or not the server185        // was actually registered on this connection. The post-condition (server186        // not registered here) holds either way, and a duplicate/retried unregister187        // must not surface as an error — the WS client only needs to know it's gone.188        return { kind: 'unregistered', server };189    }190    handleMessage(id, payload) {191        if (typeof id !== 'string' || id.length === 0) {192            return {193                kind: 'error',194                code: 'invalid_id',195                message: '`id` must be a non-empty string',196            };197        }198        if (payload === null || typeof payload !== 'object') {199            return {200                kind: 'error',201                code: 'invalid_payload',202                message: '`payload` must be a JSON-RPC message object',203            };204        }205        const resolved = this.registrar.resolveMessage(id, payload);206        return resolved207            ? { kind: 'message_resolved', id }208            : { kind: 'ignored', reason: `no pending request for id '${id}'` };209    }210    /** Currently-registered server names on this connection. */211    registeredServers() {212        return this.registrar.registeredServers();213    }214    /** In-flight `mcp_message` round-trip count (for tests / accounting). */215    pendingCount() {216        return this.registrar.pendingCount();217    }218    /**219     * Tear the connection down: reject pending, forget servers, and best-effort220     * remove each from the provider. Idempotent.221     */222    async dispose(reason = 'client MCP WS connection closed') {223        if (this.disposed)224            return;225        this.disposed = true;226        const servers = this.registrar.registeredServers();227        this.registrar.close(reason);228        if (this.provider) {229            await Promise.allSettled(servers.map((server) => this.provider.unregisterClientMcpServer(server)));230        }231    }232}233//# sourceMappingURL=client-mcp-ws.js.map