@xalia/agent 0.5.4 → 0.5.5

package/src/chat/conversationManager.ts

@@ -15,11 +15,17 @@ import { createAgentWithSkills } from "../agent/agentUtils";
 import type {
   ClientToServer,
   ServerToClient,
-  ClientUserMessage,
+  ServerMcpServerAdded,
+  ServerMcpServerRemoved,
+  ServerMcpServerToolEnabled,
+  ServerMcpServerToolDisabled,
+  ServerSystemPromptUpdated,
+  ServerModelUpdated,
 } from "./messages";
 import { AsyncQueue } from "./asyncQueue";
 import { Database, UserData } from "./db";
 import { AsyncLock } from "../utils/asyncLock";
+import { McpServerInfo, McpServerManager } from "../agent/mcpServerManager";
 
 const logger = getLogger();
 
@@ -38,20 +44,29 @@ type QueuedClientMessage<T extends ClientToServer = ClientToServer> = {
  * Describes a Session (conversation) with connected participants.
  */
 export class OpenSession {
+  public db: Database;
   public onEmpty: () => void;
+  /// Map of user identifier to connection
   public connections: Record<string, ws.WebSocket>;
   public agent: Agent;
+  public sessionUUID: string;
+  public agentProfileUUID: string;
   public skillManager: SkillManager;
   public messageQueue: AsyncQueue<QueuedClientMessage>;
   public curAgentMsgId: string | undefined;
 
   constructor(
+    db: Database,
     agent: Agent,
+    sessionUUID: string,
+    agentProfileUUID: string,
     sudoMcpServerManager: SkillManager,
     onEmpty: () => void
   ) {
-    // public agent: Agent,
+    this.db = db;
     this.agent = agent;
+    this.sessionUUID = sessionUUID;
+    this.agentProfileUUID = agentProfileUUID;
     this.skillManager = sudoMcpServerManager;
     this.onEmpty = onEmpty;
     this.connections = {};
@@ -71,6 +86,25 @@ export class OpenSession {
     this.broadcast({ type: "user_joined", user: userName });
     this.connections[userName] = ws;
 
+    // TODO:
+    //
+    //   - send conversation state (ServerHistory)
+
+    const briefs = this.skillManager.getServerBriefs();
+    this.sendTo(userName, { type: "mcp_server_briefs", server_briefs: briefs });
+
+    const msm = this.agent.getMcpServerManager();
+    const mcpServerNames = msm.getMcpServerNames();
+    for (const serverName of mcpServerNames) {
+      const info = msm.getMcpServer(serverName);
+      this.sendTo(userName, {
+        type: "mcp_server_added",
+        server_name: serverName,
+        tools: info.getTools(),
+        enabled_tools: Object.keys(info.getEnabledTools()),
+      });
+    }
+
     ws.on("message", async (message: ws.RawData) => {
       logger.debug(`[convMgr]: got message: (from ${userName}): ${message}`);
       const msgStr = message.toString();
@@ -102,14 +136,71 @@ export class OpenSession {
         `${JSON.stringify(queuedMessage.msg)}`
     );
 
-    const msg = queuedMessage.msg;
-    switch (msg.type) {
-      case "msg":
-        return this.onChatMessage(
-          queuedMessage as QueuedClientMessage<ClientUserMessage>
-        );
-      default:
-        throw `unknown message: ${JSON.stringify(queuedMessage)}`;
+    // In general, handlers return a message to be broadcast.  Errors are
+    // handled by returning an error to just the sender.  Handlers can
+    // also broadcast and send directly.
+
+    try {
+      const msg = queuedMessage.msg;
+      let broadcastMsg: ServerToClient[] | ServerToClient | undefined =
+        undefined;
+      switch (msg.type) {
+        case "msg":
+          broadcastMsg = await this.onChatMessage(
+            msg.message,
+            queuedMessage.from
+          );
+          break;
+        case "add_mcp_server":
+          broadcastMsg = await this.onAddMcpServer(
+            msg.server_name,
+            msg.enable_all
+          );
+          break;
+        case "remove_mcp_server":
+          broadcastMsg = await this.onRemoveMcpServer(msg.server_name);
+          break;
+        case "enable_mcp_server_tool":
+          broadcastMsg = await this.onEnableMcpServerTool(
+            msg.server_name,
+            msg.tool
+          );
+          break;
+        case "disable_mcp_server_tool":
+          broadcastMsg = await this.onDisableMcpServerTool(
+            msg.server_name,
+            msg.tool
+          );
+          break;
+        case "enable_all_mcp_server_tools":
+          broadcastMsg = await this.onEnableAllMcpServerTools(msg.server_name);
+          break;
+        case "disable_all_mcp_server_tools":
+          broadcastMsg = await this.onDisableAllMcpServerTools(msg.server_name);
+          break;
+        case "set_system_prompt":
+          broadcastMsg = await this.onSetSystemPrompt(msg.system_prompt);
+          break;
+        case "set_model":
+          broadcastMsg = await this.onSetModel(msg.model);
+          break;
+        default:
+          throw `unknown message: ${JSON.stringify(queuedMessage)}`;
+      }
+
+      if (broadcastMsg) {
+        if (broadcastMsg instanceof Array) {
+          broadcastMsg.map((msg) => this.broadcast(msg));
+        } else {
+          this.broadcast(broadcastMsg);
+        }
+      }
+    } catch (err: unknown) {
+      if (typeof err === "string") {
+        this.sendTo(queuedMessage.from, { type: "error", message: err });
+      } else {
+        throw err;
+      }
     }
   }
 
@@ -124,24 +215,163 @@ export class OpenSession {
     });
   }
 
-  async onChatMessage(
-    queuedMessage: QueuedClientMessage<ClientUserMessage>
-  ): Promise<void> {
-    const msg = queuedMessage.msg;
-    const userToken = queuedMessage.from;
+  async onChatMessage(message: string, userToken: string): Promise<undefined> {
+    // We manually broadcast the user's message here and start the agent
+    // conversation, and then wait to get back all data from the agent before
+    // processing further messages from clients.
+
     const msgId = uuidv4();
     this.broadcast({
       type: "user_msg",
       message_id: msgId,
-      message: msg.message,
+      message,
       from: userToken,
     });
 
-    // Messages will be handled by the Agent.onMessage callback.  We await the
-    // response here before processing further messages.
-
     this.curAgentMsgId = `${msgId}-resp`;
-    await this.agent.userMessage(msg.message, undefined, userToken);
+    await this.agent.userMessageEx(message, undefined, userToken);
+  }
+
+  async onAddMcpServer(
+    serverName: string,
+    enableAll: boolean
+  ): Promise<ServerMcpServerAdded> {
+    logger.info(
+      `[onAddMcpServer]: Adding server ${serverName} (enable_all: ${enableAll})`
+    );
+
+    const mcpServerManager = this.agent.getMcpServerManager();
+    if (mcpServerManager.hasMcpServer(serverName)) {
+      throw `${serverName} already added`;
+    }
+
+    if (!this.skillManager.hasServer(serverName)) {
+      throw `no such server: ${serverName}`;
+    }
+
+    await this.skillManager.addMcpServer(serverName, enableAll);
+    mcpServerManager.enableAllTools(serverName);
+
+    // Save changes to the AgentProfile
+
+    await this.updateAgentProfile();
+
+    // Broadcast the message to all participants.
+
+    const server = mcpServerManager.getMcpServer(serverName);
+    const tools = server.getTools();
+    const enabled_tools = Object.keys(server.getEnabledTools());
+    return {
+      type: "mcp_server_added",
+      server_name: serverName,
+      tools,
+      enabled_tools,
+    };
+  }
+
+  async onRemoveMcpServer(
+    server_name: string
+  ): Promise<ServerMcpServerRemoved> {
+    logger.info(`[onRemoveMcpServer]: Removing server ${server_name}`);
+
+    const mcpServerManager = this.agent.getMcpServerManager();
+    if (!mcpServerManager.hasMcpServer(server_name)) {
+      throw `${server_name} not enabled`;
+    }
+
+    await mcpServerManager.removeMcpServer(server_name);
+    await this.updateAgentProfile();
+
+    return {
+      type: "mcp_server_removed",
+      server_name,
+    };
+  }
+
+  async onEnableMcpServerTool(
+    server_name: string,
+    tool: string
+  ): Promise<ServerMcpServerToolEnabled> {
+    const msm = this.agent.getMcpServerManager();
+    this.ensureMcpServerAndTool(msm, server_name, tool);
+    msm.enableTool(server_name, tool);
+
+    await this.updateAgentProfile();
+
+    return { type: "mcp_server_tool_enabled", server_name, tool };
+  }
+
+  async onDisableMcpServerTool(
+    server_name: string,
+    tool: string
+  ): Promise<ServerMcpServerToolDisabled> {
+    const msm = this.agent.getMcpServerManager();
+    this.ensureMcpServerAndTool(msm, server_name, tool);
+    msm.disableTool(server_name, tool);
+
+    await this.updateAgentProfile();
+
+    return { type: "mcp_server_tool_disabled", server_name, tool };
+  }
+
+  async onEnableAllMcpServerTools(
+    server_name: string
+  ): Promise<ServerMcpServerToolEnabled[]> {
+    // We reimplement the logic to enable any disabled tools so we can
+    // construct messages along the way.
+
+    const msm = this.agent.getMcpServerManager();
+    const server = this.ensureMcpServer(msm, server_name);
+    const enabledTools = server.getEnabledTools();
+    const msgs: ServerMcpServerToolEnabled[] = [];
+    for (const tool of server.getTools()) {
+      if (!enabledTools[tool.name]) {
+        msm.enableTool(server_name, tool.name);
+        msgs.push({
+          type: "mcp_server_tool_enabled",
+          server_name,
+          tool: tool.name,
+        });
+      }
+    }
+
+    await this.updateAgentProfile();
+
+    return msgs;
+  }
+
+  async onDisableAllMcpServerTools(
+    server_name: string
+  ): Promise<ServerMcpServerToolDisabled[]> {
+    // We reimplement the logic to disable all enabled tools so we can
+    // construct messages along the way.
+
+    const msm = this.agent.getMcpServerManager();
+    const server = this.ensureMcpServer(msm, server_name);
+    const enabledTools = server.getEnabledTools();
+    const msgs: ServerMcpServerToolDisabled[] = [];
+    for (const tool in enabledTools) {
+      msm.disableTool(server_name, tool);
+      msgs.push({ type: "mcp_server_tool_disabled", server_name, tool });
+    }
+
+    await this.updateAgentProfile();
+
+    return msgs;
+  }
+
+  async onSetSystemPrompt(
+    system_prompt: string
+  ): Promise<ServerSystemPromptUpdated> {
+    this.agent.setSystemPrompt(system_prompt);
+    await this.updateAgentProfile();
+    return { type: "system_prompt_updated", system_prompt };
+  }
+
+  async onSetModel(model: string): Promise<ServerModelUpdated> {
+    this.agent.setModel(model);
+    await this.updateAgentProfile();
+    return { type: "model_updated", model };
   }
 
   broadcast(msg: ServerToClient) {
@@ -151,6 +381,47 @@ export class OpenSession {
       ws.send(msgString);
     }
   }
+
+  sendTo(userName: string, msg: ServerToClient) {
+    const ws = this.connections[userName];
+    const msgString = JSON.stringify(msg);
+    logger.info(`[sendTo]: (${userName}) msg: ${msgString}`);
+    assert(ws);
+    ws.send(msgString);
+  }
+
+  private ensureMcpServer(
+    msm: McpServerManager,
+    serverName: string
+  ): McpServerInfo {
+    const server = msm.getMcpServer(serverName);
+    if (!server) {
+      throw `${serverName} not added`;
+    }
+    return server;
+  }
+
+  private ensureMcpServerAndTool(
+    msm: McpServerManager,
+    serverName: string,
+    toolName: string
+  ) {
+    const server = this.ensureMcpServer(msm, serverName);
+    const tool = server.getTool(toolName);
+    if (!tool) {
+      throw `Tool ${toolName} on ${serverName} not found`;
+    }
+    return tool;
+  }
+
+  private async updateAgentProfile(): Promise<void> {
+    const profile = this.agent.getAgentProfile();
+    logger.debug(
+      `[updateAgentProfile]: uuid: ${this.agentProfileUUID} profile: ` +
+        JSON.stringify(profile)
+    );
+    return this.db.updateAgentProfile(this.agentProfileUUID, profile);
+  }
 }
 
 /**
@@ -184,11 +455,12 @@ export class ConversationManager {
         sessionId,
         llmApiKey,
         xmcpApiKey,
-        userData.nickname || userData.user_uuid,
+        userData.nickname || userData.uuid,
         ws
       );
     });
   }
+
   /**
    * Must be called while holding the openSessionsLock
    */
@@ -223,29 +495,56 @@ export class ConversationManager {
     const conversation =
       sessionData.conversation as unknown as ChatCompletionMessageParam[];
 
-    const agentProfile = await this.db.getAgentProfileById(
-      sessionData.agent_profile_uuid
-    );
+    const agentProfileUUID = sessionData.agent_profile_uuid;
+    const agentProfile = await this.db.getAgentProfileById(agentProfileUUID);
     if (!agentProfile) {
-      throw `no such agent profile ${sessionData.agent_profile_uuid}`;
+      throw `no such agent profile ${agentProfileUUID}`;
     }
 
+    // TODO: store some owner data on the OpenSession object itself?
+
+    const owner = await this.db.getUserFromUuid(sessionData.user_uuid);
+    assert(owner, `no owner for session ${JSON.stringify(sessionData)}`);
+    const ownerUserName = owner.nickname;
+    if (!ownerUserName) {
+      throw (
+        `user ${sessionData.user_uuid} has no user name - ` +
+        "cannot create chat session"
+      );
+    }
+
+    // Access to the OpenSession (once it is iniailized
+    const context: { openSession?: OpenSession } = {};
+
     const platform = {
       openUrl: (
-        _url: string,
+        url: string,
         _authResultP: Promise<boolean>,
-        _displayName: string
+        display_name: string
       ) => {
-        throw "unimpl";
+        // These requests are always passed to the original owner, since it is
+        // his settings that will be used for all mcp servers.
+
+        if (context.openSession) {
+          const conn = openSession.connections[ownerUserName];
+          if (conn) {
+            openSession.sendTo(ownerUserName, {
+              type: "open_url",
+              url,
+              display_name,
+            });
+          } else {
+            throw `user ${ownerUserName} must authenticate`;
+          }
+        } else {
+          throw `no open session ${sessionData.uuid}`;
+        }
       },
       load: (_filename: string): Promise<string> => {
-        throw "unimpl";
+        throw "unimpl platform.load";
       },
     };
 
-    // Forward agent messages to the OpenSession (once it is iniailized
-
-    const context: { openSession?: OpenSession } = {};
     const onMessage = async (msg: string, end: boolean) => {
       logger.debug(`[onMessage] msg: ${msg}, end: ${end}`);
       assert(context.openSession);
@@ -256,6 +555,12 @@ export class ConversationManager {
       toolCall: ChatCompletionMessageToolCall
     ): Promise<boolean> => {
       logger.debug(`[onToolCall] : ${JSON.stringify(toolCall)}`);
+      assert(context.openSession);
+      context.openSession.broadcast({
+        type: "agent_tool_call",
+        message_id: toolCall.id,
+        message: toolCall,
+      });
       return true;
     };
 
@@ -274,7 +579,14 @@ export class ConversationManager {
       true
     );
 
-    openSession = new OpenSession(agent, smsm, onEmpty);
+    openSession = new OpenSession(
+      this.db,
+      agent,
+      sessionId,
+      agentProfileUUID,
+      smsm,
+      onEmpty
+    );
     context.openSession = openSession;
     this.openSessions[sessionId] = openSession;
     openSession.join(userName, ws);

package/src/chat/db.ts

@@ -31,7 +31,7 @@ export function resolveCompoundName(name: string): string | [string, string] {
 }
 
 export type UserData = {
-  user_uuid: supabase.Tables<"api_keys">["user_uuid"];
+  uuid: supabase.Tables<"api_keys">["user_uuid"];
   nickname: supabase.Tables<"users">["nickname"];
 };
 
@@ -67,11 +67,24 @@ export class Database {
     }
 
     return {
-      user_uuid: data.user_uuid,
+      uuid: data.user_uuid,
       nickname: data.users.nickname || `user ${data.user_uuid}`,
     };
   }
 
+  async getUserFromUuid(user_uuid: string): Promise<UserData | undefined> {
+    const { data, error } = await this.client
+      .from("users")
+      .select("*")
+      .eq("uuid", user_uuid)
+      .maybeSingle();
+    if (error) {
+      throw error;
+    }
+
+    return data;
+  }
+
   async createUser(
     user_uuid: string,
     email: string,
@@ -196,6 +209,19 @@ export class Database {
     return data[0].uuid;
   }
 
+  async updateAgentProfile(uuid: string, profile: AgentProfile): Promise<void> {
+    const payload: supabase.TablesUpdate<"agent_profiles"> = {
+      profile: profile as unknown as supabase.Json,
+    };
+    const { error } = await this.client
+      .from("agent_profiles")
+      .update(payload)
+      .eq("uuid", uuid);
+    if (error) {
+      throw error;
+    }
+  }
+
   async clearAgentProfiles(): Promise<void> {
     await this.client.from("agent_profiles").delete().neq("uuid", "");
   }

package/src/chat/frontendClient.ts

@@ -0,0 +1,123 @@
+import { ServerToClient, ClientToServer } from "./messages";
+
+type OnMessageCallback = (msg: ServerToClient) => void;
+type OnConnectionClosedCallback = () => void;
+
+export class FrontendChatClient {
+  private ws: WebSocket;
+  private onMessageCB: OnMessageCallback;
+  private onConnectionClosedCB: OnConnectionClosedCallback;
+  private closed: boolean;
+
+  private constructor(
+    ws: WebSocket,
+    onMessageCB: OnMessageCallback,
+    onConnectionClosedCB: OnConnectionClosedCallback
+  ) {
+    this.ws = ws;
+    this.onMessageCB = onMessageCB;
+    this.onConnectionClosedCB = onConnectionClosedCB;
+    this.closed = false;
+  }
+
+  static async initWithParams(
+    host: string,
+    port: number,
+    token: string,
+    params: Record<string, string>,
+    onMessageCB: OnMessageCallback,
+    onConnectionClosedCB: OnConnectionClosedCallback
+  ): Promise<FrontendChatClient> {
+    return new Promise((resolve, reject) => {
+      const urlParams = new URLSearchParams(params);
+      const url = `ws://${host}:${port}?${urlParams}`;
+
+      const ws = new WebSocket(url, [token]);
+      console.log("created ws");
+
+      const client = new FrontendChatClient(
+        ws,
+        onMessageCB,
+        onConnectionClosedCB
+      );
+
+      ws.onopen = async () => {
+        console.log("opened");
+
+        ws.onmessage = (ev: MessageEvent) => {
+          try {
+            const msgData = ev.data;
+            if (typeof msgData !== "string") {
+              throw `expected "string" data, got ${typeof msgData}`;
+            }
+            console.debug(`[client.onmessage]: ${msgData}`);
+            const msg: ServerToClient = JSON.parse(msgData);
+            client.onMessageCB(msg);
+          } catch (e) {
+            client.close();
+            throw e;
+          }
+        };
+        resolve(client);
+      };
+
+      ws.onclose = (event) => {
+        console.log("closed");
+        console.log(
+          `[client] WebSocket connection closed: ${JSON.stringify(event)}`
+        );
+        client.closed = true;
+        onConnectionClosedCB();
+      };
+
+      ws.onerror = (error) => {
+        console.error("[client] WebSocket error:", error);
+        reject(error);
+
+        client.closed = true;
+        onConnectionClosedCB();
+      };
+    });
+  }
+
+  public static async init(
+    host: string,
+    port: number,
+    token: string,
+    onMessageCB: OnMessageCallback,
+    onConnectionClosedCB: OnConnectionClosedCallback,
+    sessionId: string = "untitled",
+    agentProfileId: string | undefined = undefined
+  ): Promise<FrontendChatClient> {
+    const params: Record<string, string> = { session_id: sessionId };
+    if (agentProfileId) {
+      params["agent_profile_id"] = agentProfileId;
+    }
+    return FrontendChatClient.initWithParams(
+      host,
+      port,
+      token,
+      params,
+      onMessageCB,
+      onConnectionClosedCB
+    );
+  }
+
+  public sendMessage(message: ClientToServer): void {
+    if (this.closed) {
+      throw new Error("Cannot send message on closed connection");
+    }
+    const data = JSON.stringify(message);
+    this.ws.send(data);
+  }
+
+  public close(): void {
+    this.closed = true;
+    this.onConnectionClosedCB();
+    this.ws.close();
+  }
+
+  public isClosed(): boolean {
+    return this.closed;
+  }
+}