@xalia/agent 0.5.4 → 0.5.5

package/src/chat/client.ts

@@ -1,32 +1,230 @@
-import * as dotenv from "dotenv";
-import { getLogger } from "@xalia/xmcp/sdk";
+import { Tool } from "@modelcontextprotocol/sdk/types.js";
 import { strict as assert } from "assert";
 import * as websocket from "ws";
-import { ServerToClient, ClientToServer } from "./messages";
 
-dotenv.config();
+import { AgentProfile, McpServerBrief, getLogger } from "@xalia/xmcp/sdk";
+
+import { ISkillManager } from "../agent/sudoMcpServerManager";
+import {
+  IMcpServerManager,
+  McpServerInfo,
+  McpServerInfoRW,
+} from "../agent/mcpServerManager";
+
+import { ServerToClient, ClientToServer } from "./messages";
+import { ChatCompletionMessageParam, IConversation } from "../agent/agent";
+import { IPlatform } from "../agent/iplatform";
 
 const logger = getLogger();
 
 type OnMessageCallback = (msg: ServerToClient) => void;
 
-type OnConnectionClosedCallback = () => void;
+type OnConnectionClosedCallback = () => Promise<void>;
+
+interface IMessageSender {
+  sendMessage(message: ClientToServer): void;
+}
+
+class RemoteMcpServerManager implements IMcpServerManager {
+  private mcpServers: { [serverName: string]: McpServerInfoRW } = {};
+
+  constructor(private sender: IMessageSender) {
+    this.mcpServers = {};
+  }
+
+  hasMcpServer(mcpServerName: string): boolean {
+    return !!this.mcpServers[mcpServerName];
+  }
+
+  getMcpServerNames(): string[] {
+    return Object.keys(this.mcpServers);
+  }
+
+  getMcpServer(mcpServerName: string): McpServerInfo {
+    const server = this.mcpServers[mcpServerName];
+    if (server) {
+      return server;
+    }
+    throw Error(`unknown server ${mcpServerName}`);
+  }
+
+  async removeMcpServer(mcpServerName: string): Promise<void> {
+    if (!this.mcpServers[mcpServerName]) {
+      throw Error(`no server ${mcpServerName} (removeMcpServer)`);
+    }
+
+    this.sender.sendMessage({
+      type: "remove_mcp_server",
+      server_name: mcpServerName,
+    });
+  }
+
+  enableAllTools(mcpServerName: string): void {
+    if (!this.mcpServers[mcpServerName]) {
+      throw Error(`no server ${mcpServerName} (enableAllTools)`);
+    }
+
+    this.sender.sendMessage({
+      type: "enable_all_mcp_server_tools",
+      server_name: mcpServerName,
+    });
+  }
+
+  disableAllTools(mcpServerName: string): void {
+    if (!this.mcpServers[mcpServerName]) {
+      throw Error(`no server ${mcpServerName} (disableAllTools)`);
+    }
+
+    this.sender.sendMessage({
+      type: "disable_all_mcp_server_tools",
+      server_name: mcpServerName,
+    });
+  }
+
+  enableTool(mcpServerName: string, toolName: string): void {
+    const server = this.mcpServers[mcpServerName];
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (enableTool)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    this.sender.sendMessage({
+      type: "enable_mcp_server_tool",
+      server_name: mcpServerName,
+      tool: toolName,
+    });
+  }
 
-export class ChatClient {
+  disableTool(mcpServerName: string, toolName: string): void {
+    const server = this.mcpServers[mcpServerName];
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (disableTool)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    this.sender.sendMessage({
+      type: "disable_mcp_server_tool",
+      server_name: mcpServerName,
+      tool: toolName,
+    });
+  }
+
+  onMcpServerAdded(
+    mcpServerName: string,
+    tools: Tool[],
+    enabled_tools: string[]
+  ) {
+    logger.debug(
+      `[onMcpServerAdded]: ${mcpServerName}, tools: ${JSON.stringify(tools)}` +
+        `, enabled: ${JSON.stringify(enabled_tools)}`
+    );
+
+    const mcpServerInfo = new McpServerInfoRW(tools);
+    for (const tool of enabled_tools) {
+      mcpServerInfo.enableTool(tool);
+    }
+    this.mcpServers[mcpServerName] = mcpServerInfo;
+  }
+
+  onMcpServerRemoved(mcpServerName: string) {
+    delete this.mcpServers[mcpServerName];
+  }
+
+  onMcpServerToolEnabled(mcpServerName: string, toolName: string) {
+    const server = this.mcpServers[mcpServerName];
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (onMcpServerToolEnabled)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    server.enableTool(toolName);
+  }
+
+  onMcpServerToolDisabled(mcpServerName: string, toolName: string) {
+    const server = this.mcpServers[mcpServerName];
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (onMcpServerToolDisabled)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    server.disableTool(toolName);
+  }
+}
+
+class RemoteSudoMcpServerManager implements ISkillManager {
+  private briefsMap: Record<string, McpServerBrief>;
+
+  constructor(
+    private sender: IMessageSender,
+    private briefs: McpServerBrief[],
+    private msm: RemoteMcpServerManager
+  ) {
+    this.briefsMap = {};
+    briefs.forEach((b) => {
+      this.briefsMap[b.name] = b;
+    });
+  }
+
+  getMcpServerManager(): IMcpServerManager {
+    return this.msm;
+  }
+
+  getServerBriefs(): McpServerBrief[] {
+    return this.briefs;
+  }
+
+  async addMcpServer(server_name: string, enable_all: boolean): Promise<void> {
+    if (!this.briefsMap[server_name]) {
+      throw Error(`no such server ${server_name} (addMcpServer)`);
+    }
+
+    this.sender.sendMessage({
+      type: "add_mcp_server",
+      server_name,
+      enable_all,
+    });
+  }
+}
+
+export class ChatClient implements IMessageSender, IConversation {
+  private platform: IPlatform;
   private ws: websocket.WebSocket;
   private onMessageCB: OnMessageCallback;
   private onConnectionClosedCB: OnConnectionClosedCallback;
   private closed: boolean;
+  private msm: RemoteMcpServerManager;
+  private smsm: RemoteSudoMcpServerManager;
+  private systemPrompt: string;
+  private model: string;
 
   private constructor(
+    platform: IPlatform,
     ws: websocket.WebSocket,
     onMessageCB: OnMessageCallback,
-    onConnectionClosedCB: OnConnectionClosedCallback
+    onConnectionClosedCB: OnConnectionClosedCallback,
+    serverBriefs: McpServerBrief[]
   ) {
+    this.platform = platform;
     this.ws = ws;
     this.onMessageCB = onMessageCB;
     this.onConnectionClosedCB = onConnectionClosedCB;
     this.closed = false;
+    this.msm = new RemoteMcpServerManager(this);
+    this.smsm = new RemoteSudoMcpServerManager(this, serverBriefs, this.msm);
+    this.systemPrompt = "";
+    this.model = "";
   }
 
   static async initWithParams(
@@ -35,16 +233,40 @@ export class ChatClient {
     token: string,
     params: Record<string, string>,
     onMessageCB: OnMessageCallback,
-    onConnectionClosedCB: OnConnectionClosedCallback
+    onConnectionClosedCB: OnConnectionClosedCallback,
+    platform: IPlatform
   ): Promise<ChatClient> {
-    return new Promise((r, e) => {
+    return new Promise((resolveClient, e) => {
       const urlParams = new URLSearchParams(params);
       const url = `ws://${host}:${port}?${urlParams}`;
 
       const ws = new websocket.WebSocket(url, [token]);
       logger.info("created ws");
 
-      const client = new ChatClient(ws, onMessageCB, onConnectionClosedCB);
+      let client: ChatClient | undefined = undefined;
+      const onMsg = (msg: ServerToClient) => {
+        if (msg.type === "mcp_server_briefs") {
+          // This should be received once, as the first message at startup.
+          assert(!client);
+          client = new ChatClient(
+            platform,
+            ws,
+            onMessageCB,
+            onConnectionClosedCB,
+            msg.server_briefs
+          );
+          resolveClient(client);
+        } else {
+          assert(client);
+
+          // Pass all messages to our internal handler (to update mcp state,
+          // etc) before sending to the client code.
+
+          // logger.debug(`[ChatClient.init(onMsg)]: ${JSON.stringify(msg)}`);
+          client.handleMessageInternal(msg);
+          client.onMessageCB(msg);
+        }
+      };
 
       ws.onopen = async () => {
         logger.info("opened");
@@ -56,59 +278,46 @@ export class ChatClient {
               throw `expected "string" data, got ${typeof msgData}`;
             }
             logger.debug(`[client.onmessage]: ${msgData}`);
-            const msg: ServerToClient = JSON.parse(msgData);
-            client.onMessageCB(msg);
+            onMsg(JSON.parse(msgData) as ServerToClient);
           } catch (e) {
-            client.close();
+            if (client) {
+              client.close();
+            }
             throw e;
           }
         };
-        r(client);
       };
 
-      ws.onclose = (err) => {
+      ws.onclose = async (err) => {
         logger.info("closed");
         logger.info(
           `[client] WebSocket connection closed: ${JSON.stringify(err)}`
         );
-        client.closed = true;
-        onConnectionClosedCB();
+        if (client) {
+          client.closed = true;
+        }
+        await onConnectionClosedCB();
       };
 
-      ws.onerror = (error) => {
+      ws.onerror = async (error) => {
         logger.error("[client] WebSocket error:", JSON.stringify(error));
         e(error);
 
-        client.closed = true;
-        onConnectionClosedCB();
+        if (client) {
+          client.closed = true;
+        }
+        await onConnectionClosedCB();
       };
     });
   }
 
-  // public static async initWithSession(
-  //   host: string,
-  //   port: number,
-  //   token: string,
-  //   sessionId: string,
-  //   onMessageCB: OnMessageCallback,
-  //   onConnectionClosedCB: OnConnectionClosedCallback
-  // ): Promise<ChatClient> {
-  //   return ChatClient.initWithParams(
-  //     host,
-  //     port,
-  //     token,
-  //     { session_id: sessionId },
-  //     onMessageCB,
-  //     onConnectionClosedCB
-  //   );
-  // }
-
   public static async init(
     host: string,
     port: number,
     token: string,
     onMessageCB: OnMessageCallback,
     onConnectionClosedCB: OnConnectionClosedCallback,
+    platform: IPlatform,
     sessionId: string = "untitled",
     agentProfileId: string | undefined = undefined
   ): Promise<ChatClient> {
@@ -122,13 +331,60 @@ export class ChatClient {
       token,
       params,
       onMessageCB,
-      onConnectionClosedCB
+      onConnectionClosedCB,
+      platform
     );
   }
 
-  public sendMessage(message: ClientToServer): void {
+  public getSudoMcpServerManager(): ISkillManager {
+    return this.smsm;
+  }
+
+  public getConversation(): ChatCompletionMessageParam[] {
+    throw "unimpl";
+  }
+
+  public getAgentProfile(): AgentProfile {
+    throw "unimpl";
+  }
+
+  getSystemPrompt(): string {
+    return this.systemPrompt;
+  }
+
+  setSystemPrompt(system_prompt: string): void {
+    // Don't set system prompt here.  Wait until we get confirmation from the
+    // server.
+    this.sendMessage({ type: "set_system_prompt", system_prompt });
+  }
+
+  getModel(): string {
+    return this.model;
+  }
+
+  setModel(model: string): void {
+    // Don't set model here.  Wait until we get confirmation from the server.
+    this.sendMessage({ type: "set_model", model });
+  }
+
+  userMessage(msg?: string, imageB64?: string): void {
+    assert(msg);
+    assert(!imageB64, "images not supported in Chat yet");
+
+    this.sendMessage({
+      type: "msg",
+      message: msg,
+    });
+  }
+
+  resetConversation(): void {
+    throw "resetConversation not implemented for ChatClient";
+  }
+
+  public sendMessage(message: ClientToServer): undefined {
     assert(!this.closed);
     const data = JSON.stringify(message);
+    logger.debug(`[ChatClient.sendMessage] ${data}`);
     this.ws.send(data);
   }
 
@@ -137,6 +393,63 @@ export class ChatClient {
     this.onConnectionClosedCB();
     this.ws.close();
   }
+
+  handleMessageInternal(message: ServerToClient): void {
+    switch (message.type) {
+      //
+      // State updates
+      //
+
+      case "mcp_server_added":
+        this.msm.onMcpServerAdded(
+          message.server_name,
+          message.tools,
+          message.enabled_tools
+        );
+        break;
+      case "mcp_server_removed":
+        this.msm.onMcpServerRemoved(message.server_name);
+        break;
+      case "mcp_server_tool_enabled":
+        this.msm.onMcpServerToolEnabled(message.server_name, message.tool);
+        break;
+      case "mcp_server_tool_disabled":
+        this.msm.onMcpServerToolDisabled(message.server_name, message.tool);
+        break;
+      case "system_prompt_updated":
+        this.systemPrompt = message.system_prompt;
+        break;
+      case "model_updated":
+        this.model = message.model;
+        break;
+
+      //
+      // Actions
+      //
+
+      case "open_url":
+        this.platform.openUrl(
+          message.url,
+          new Promise(() => {}), // TODO: Why do we need this?  Remove it.
+          message.display_name
+        );
+        break;
+      case "approve_tool_call":
+        throw "unimpl approve_tool_call";
+      // break;
+
+      //
+      // Ignore other messages - the owner (the UI layer) can handle them at
+      // its discretion.
+      //
+
+      default:
+        logger.debug(
+          `[handleMessageInternal]: ignoring message: ${message.type}`
+        );
+        break;
+    }
+  }
 }
 
 // TODO: remove this