claude-code/packages/@ant/claude-for-chrome-mcp/src/mcpSocketPool.ts

import {
  createMcpSocketClient,
  SocketConnectionError,
} from "./mcpSocketClient.js";
import type { McpSocketClient } from "./mcpSocketClient.js";
import type {
  ClaudeForChromeContext,
  PermissionMode,
  PermissionOverrides,
} from "./types.js";

/**
 * Manages connections to multiple Chrome native host sockets (one per Chrome profile).
 * Routes tool calls to the correct socket based on tab ID.
 *
 * For `tabs_context_mcp`: queries all connected sockets and merges results.
 * For other tools: routes based on the `tabId` argument using a routing table
 * built from tabs_context_mcp responses.
 */
export class McpSocketPool {
  private clients: Map<string, McpSocketClient> = new Map();
  private tabRoutes: Map<number, string> = new Map();
  private context: ClaudeForChromeContext;
  private notificationHandler:
    | ((notification: { method: string; params?: Record<string, unknown> }) => void)
    | null = null;

  constructor(context: ClaudeForChromeContext) {
    this.context = context;
  }

  public setNotificationHandler(
    handler: (notification: {
      method: string;
      params?: Record<string, unknown>;
    }) => void,
  ): void {
    this.notificationHandler = handler;
    for (const client of this.clients.values()) {
      client.setNotificationHandler(handler);
    }
  }

  /**
   * Discover available sockets and ensure at least one is connected.
   */
  public async ensureConnected(): Promise<boolean> {
    const { logger, serverName } = this.context;

    this.refreshClients();

    // Try to connect any disconnected clients
    const connectPromises: Promise<boolean>[] = [];
    for (const client of this.clients.values()) {
      if (!client.isConnected()) {
        connectPromises.push(
          client.ensureConnected().catch(() => false),
        );
      }
    }

    if (connectPromises.length > 0) {
      await Promise.all(connectPromises);
    }

    const connectedCount = this.getConnectedClients().length;
    if (connectedCount === 0) {
      logger.info(`[${serverName}] No connected sockets in pool`);
      return false;
    }

    logger.info(`[${serverName}] Socket pool: ${connectedCount} connected`);
    return true;
  }

  /**
   * Call a tool, routing to the correct socket based on tab ID.
   * For tabs_context_mcp, queries all sockets and merges results.
   */
  public async callTool(
    name: string,
    args: Record<string, unknown>,
    _permissionOverrides?: PermissionOverrides,
  ): Promise<unknown> {
    if (name === "tabs_context_mcp") {
      return this.callTabsContext(args);
    }

    // Route by tabId if present
    const tabId = args.tabId as number | undefined;
    if (tabId !== undefined) {
      const socketPath = this.tabRoutes.get(tabId);
      if (socketPath) {
        const client = this.clients.get(socketPath);
        if (client?.isConnected()) {
          return client.callTool(name, args);
        }
      }
      // Tab route not found or client disconnected — fall through to any connected
    }

    // Fallback: use first connected client
    const connected = this.getConnectedClients();
    if (connected.length === 0) {
      throw new SocketConnectionError(
        `[${this.context.serverName}] No connected sockets available`,
      );
    }
    return connected[0]!.callTool(name, args);
  }

  public async setPermissionMode(
    mode: PermissionMode,
    allowedDomains?: string[],
  ): Promise<void> {
    const connected = this.getConnectedClients();
    await Promise.all(
      connected.map((client) => client.setPermissionMode(mode, allowedDomains)),
    );
  }

  public isConnected(): boolean {
    return this.getConnectedClients().length > 0;
  }

  public disconnect(): void {
    for (const client of this.clients.values()) {
      client.disconnect();
    }
    this.clients.clear();
    this.tabRoutes.clear();
  }

  private getConnectedClients(): McpSocketClient[] {
    return [...this.clients.values()].filter((c) => c.isConnected());
  }

  /**
   * Query all connected sockets for tabs and merge results.
   * Updates the tab routing table.
   */
  private async callTabsContext(
    args: Record<string, unknown>,
  ): Promise<unknown> {
    const { logger, serverName } = this.context;
    const connected = this.getConnectedClients();

    if (connected.length === 0) {
      throw new SocketConnectionError(
        `[${serverName}] No connected sockets available`,
      );
    }

    // If only one client, skip merging overhead
    if (connected.length === 1) {
      const result = await connected[0]!.callTool("tabs_context_mcp", args);
      this.updateTabRoutes(result, this.getSocketPathForClient(connected[0]!));
      return result;
    }

    // Query all connected clients in parallel
    const results = await Promise.allSettled(
      connected.map(async (client) => {
        const result = await client.callTool("tabs_context_mcp", args);
        const socketPath = this.getSocketPathForClient(client);
        return { result, socketPath };
      }),
    );

    // Merge tab results
    const mergedTabs: unknown[] = [];
    this.tabRoutes.clear();

    for (const settledResult of results) {
      if (settledResult.status !== "fulfilled") {
        logger.info(
          `[${serverName}] tabs_context_mcp failed on one socket: ${settledResult.reason}`,
        );
        continue;
      }

      const { result, socketPath } = settledResult.value;
      this.updateTabRoutes(result, socketPath);

      const tabs = this.extractTabs(result);
      if (tabs) {
        mergedTabs.push(...tabs);
      }
    }

    // Return merged result in the same format as the extension response
    if (mergedTabs.length > 0) {
      const tabListText = mergedTabs
        .map((t) => {
          const tab = t as { tabId: number; title: string; url: string };
          return `  • tabId ${tab.tabId}: "${tab.title}" (${tab.url})`;
        })
        .join("\n");

      return {
        result: {
          content: [
            {
              type: "text",
              text: JSON.stringify({ availableTabs: mergedTabs }),
            },
            {
              type: "text",
              text: `\n\nTab Context:\n- Available tabs:\n${tabListText}`,
            },
          ],
        },
      };
    }

    // Fallback: return first successful result as-is
    for (const settledResult of results) {
      if (settledResult.status === "fulfilled") {
        return settledResult.value.result;
      }
    }

    throw new SocketConnectionError(
      `[${serverName}] All sockets failed for tabs_context_mcp`,
    );
  }

  /**
   * Extract tab objects from a tool response to update routing table.
   */
  private updateTabRoutes(result: unknown, socketPath: string): void {
    const tabs = this.extractTabs(result);
    if (!tabs) return;

    for (const tab of tabs) {
      if (typeof tab === "object" && tab !== null && "tabId" in tab) {
        const tabId = (tab as { tabId: number }).tabId;
        this.tabRoutes.set(tabId, socketPath);
      }
    }
  }

  private extractTabs(result: unknown): unknown[] | null {
    if (!result || typeof result !== "object") return null;

    // Response format: { result: { content: [{ type: "text", text: "{\"availableTabs\":[...],\"tabGroupId\":...}" }] } }
    const asResponse = result as {
      result?: { content?: Array<{ type: string; text?: string }> };
    };
    const content = asResponse.result?.content;
    if (!content || !Array.isArray(content)) return null;

    for (const item of content) {
      if (item.type === "text" && item.text) {
        try {
          const parsed = JSON.parse(item.text);
          if (Array.isArray(parsed)) return parsed;
          // Handle { availableTabs: [...] } format
          if (parsed && Array.isArray(parsed.availableTabs)) {
            return parsed.availableTabs;
          }
        } catch {
          // Not JSON, skip
        }
      }
    }
    return null;
  }

  private getSocketPathForClient(client: McpSocketClient): string {
    for (const [path, c] of this.clients.entries()) {
      if (c === client) return path;
    }
    return "";
  }

  /**
   * Scan for available sockets and create/remove clients as needed.
   */
  private refreshClients(): void {
    const socketPaths = this.getAvailableSocketPaths();
    const { logger, serverName } = this.context;

    // Add new clients for newly discovered sockets
    for (const path of socketPaths) {
      if (!this.clients.has(path)) {
        logger.info(`[${serverName}] Adding socket to pool: ${path}`);
        const clientContext: ClaudeForChromeContext = {
          ...this.context,
          socketPath: path,
          getSocketPath: undefined,
          getSocketPaths: undefined,
        };
        const client = createMcpSocketClient(clientContext);
        client.disableAutoReconnect = true;
        if (this.notificationHandler) {
          client.setNotificationHandler(this.notificationHandler);
        }
        this.clients.set(path, client);
      }
    }

    // Remove clients for sockets that no longer exist
    for (const [path, client] of this.clients.entries()) {
      if (!socketPaths.includes(path)) {
        logger.info(`[${serverName}] Removing stale socket from pool: ${path}`);
        client.disconnect();
        this.clients.delete(path);
        for (const [tabId, socketPath] of this.tabRoutes.entries()) {
          if (socketPath === path) {
            this.tabRoutes.delete(tabId);
          }
        }
      }
    }
  }

  private getAvailableSocketPaths(): string[] {
    return this.context.getSocketPaths?.() ?? [];
  }
}

export function createMcpSocketPool(
  context: ClaudeForChromeContext,
): McpSocketPool {
  return new McpSocketPool(context);
}