@xalia/agent 0.5.8 → 0.6.0

package/src/chat/server/openSession.ts

@@ -0,0 +1,1332 @@
+import { strict as assert } from "assert";
+import { Tool } from "@modelcontextprotocol/sdk/types.js";
+
+import {
+  AgentPreferences,
+  Configuration,
+  getLogger,
+  McpServerSettings,
+  prefsGetAutoApprove,
+  prefsSetAutoApprove,
+  SavedAgentProfile,
+} from "@xalia/xmcp/sdk";
+
+import {
+  Agent,
+  ChatCompletionAssistantMessageParam,
+  ChatCompletionMessageToolCall,
+  ChatCompletionToolMessageParam,
+  createUserMessage,
+} from "../../agent/agent";
+import { SkillManager } from "../../agent/sudoMcpServerManager";
+import { McpServerInfo } from "../../agent/mcpServerManager";
+import { IAgentEventHandler } from "../../agent/iAgentEventHandler";
+import {
+  createAgentWithoutSkills,
+  DEFAULT_LLM_MODEL,
+} from "../../agent/agentUtils";
+import { IPlatform } from "../../agent/iplatform";
+
+import type {
+  ServerToClient,
+  ServerMcpServerAdded,
+  ServerMcpServerRemoved,
+  ServerMcpServerToolEnabled,
+  ServerMcpServerToolDisabled,
+  ServerSystemPromptUpdated,
+  ServerModelUpdated,
+  ServerUserMessage,
+  ClientUserMessage,
+  ServerToolAutoApprovalSet,
+  ClientSessionMessage,
+  ServerSessionError,
+  ServerUserAdded,
+  ServerUserRemoved,
+  ServerSessionInfo,
+  ServerSessionUpdate,
+  ClientSetWorkspace,
+  ClientSetMcpServerConfig,
+} from "../protocol/messages";
+import { AsyncQueue } from "../utils/asyncQueue";
+import { MultiAsyncQueue } from "../utils/multiAsyncQueue";
+import {
+  SessionCheckpoint,
+  SessionMessage,
+  SessionParticipantMap,
+  SessionData,
+  TeamParticipant,
+  SessionCreateData,
+  UserMessageData,
+} from "../data/dataModels";
+import {
+  Database,
+  createSessionParticipantMap,
+  UserData,
+} from "../data/database";
+import { ApprovalManager } from "../utils/approvalManager";
+import { ChatErrorMessage, ChatFatalError } from "../protocol/errors";
+import { addDefaultChatTools } from "./tools";
+import { ChatContextManager, ICheckpointWriter } from "./chatContextManager";
+import {
+  llmUserMessageToUserMessageData,
+  MESSAGE_INDEX_START_VALUE,
+} from "./conversation";
+import { ChatSessionFileManager } from "./sessionFileManager";
+import { IUserConnectionManager } from "./connectionManager";
+import { getErrorString } from "./errorUtils";
+import { NODE_PLATFORM } from "../../tool/nodePlatform";
+import { DbMcpServerConfigs } from "../data/dbMcpServerConfigs";
+
+/**
+ * Num messages to load at conversation startup.
+ */
+export const DEFAULT_NUM_MESSGAES = 500;
+
+const logger = getLogger();
+
+type QueuedClientMessage<
+  T extends ClientSessionMessage = ClientSessionMessage,
+> = {
+  msg: T;
+  from: string;
+};
+
+/**
+ * Implementation of ICheckpointWriter that writes into the DB.
+ */
+class DBCheckpointWriter implements ICheckpointWriter {
+  constructor(
+    private db: Database,
+    private sessionId: string
+  ) {}
+
+  writeCheckpoint(checkpoint: SessionCheckpoint): Promise<void> {
+    logger.info(
+      `[DBCheckpointWriter] writing (session ${this.sessionId}):\n` +
+        JSON.stringify(checkpoint)
+    );
+    return this.db.sessionCheckpointSet(this.sessionId, checkpoint);
+  }
+}
+
+/**
+ * Describes a Session (conversation) with connected participants.
+ *
+ * We maintain 2 queues - one for all incoming messages from clients, and
+ * another for user messages (to be gathered and sent to the Agent).  This
+ * allows processing of user messages to depend on messages from clients
+ * (using a single queue would not allow this, since later messages could not
+ * be seen until user messages had been fully processed, which could block
+ * tool approvals and other interactions).
+ */
+export class OpenSession implements IAgentEventHandler, IPlatform {
+  private readonly db: Database;
+  private /* readonly */ agent: Agent;
+  private readonly sessionUUID: string;
+  private readonly teamUUID: string | undefined;
+  private readonly userUUID: string;
+  private readonly agentProfileUUID: string;
+  private readonly sessionParticipants: SessionParticipantMap;
+  private readonly agentProfilePreferences: AgentPreferences;
+  private /* readonly */ skillManager: SkillManager;
+  private readonly connectionManager: IUserConnectionManager<ServerToClient>;
+  private readonly messageQueue: AsyncQueue<QueuedClientMessage>;
+  private readonly userMessageQueue: MultiAsyncQueue<ServerUserMessage>;
+  private readonly contextManager: ChatContextManager;
+  private readonly approvalManager: ApprovalManager;
+  private readonly savedAgentProfile: SavedAgentProfile;
+  private isPersisted: boolean;
+  private sessionTitle: string;
+  private sessionUpdatedAt: string;
+
+  private constructor(
+    db: Database,
+    agent: Agent,
+    sessionData: SessionData,
+    savedAgentProfile: SavedAgentProfile,
+    isPersisted: boolean,
+    sessionParticipants: SessionParticipantMap,
+    agentProfilePreferences: AgentPreferences,
+    skillManager: SkillManager,
+    contextManager: ChatContextManager,
+    connectionManager: IUserConnectionManager<ServerToClient>,
+    approvalManager: ApprovalManager = new ApprovalManager()
+  ) {
+    this.db = db;
+    this.agent = agent;
+    this.sessionUUID = sessionData.session_uuid;
+    this.teamUUID = sessionData.team_uuid;
+    this.userUUID = sessionData.user_uuid;
+    this.agentProfileUUID = sessionData.agent_profile_uuid;
+    this.sessionParticipants = sessionParticipants;
+    this.agentProfilePreferences = agentProfilePreferences;
+    this.skillManager = skillManager;
+    this.connectionManager = connectionManager;
+    this.messageQueue = new AsyncQueue<QueuedClientMessage>((m) =>
+      this.processMessage(m)
+    );
+
+    this.userMessageQueue = new MultiAsyncQueue<ServerUserMessage>((m) =>
+      this.processUserMessages(m)
+    );
+    this.contextManager = contextManager;
+    this.approvalManager = approvalManager;
+    this.savedAgentProfile = savedAgentProfile;
+    this.isPersisted = isPersisted;
+    this.sessionTitle = sessionData.title;
+    this.sessionUpdatedAt = sessionData.updated_at;
+  }
+
+  private static async init(
+    db: Database,
+    sessionData: SessionData,
+    savedAgentProfile: SavedAgentProfile,
+    sessionMessages: SessionMessage[],
+    sessionParticipants: SessionParticipantMap,
+    ownerData: UserData,
+    ownerApiKey: string,
+    sessionCheckpoint: SessionCheckpoint | undefined,
+    llmUrl: string,
+    xmcpUrl: string,
+    connectionManager: IUserConnectionManager<ServerToClient>
+  ): Promise<OpenSession> {
+    const sessionId = sessionData.session_uuid;
+    const fileManager = new ChatSessionFileManager(sessionId, db);
+    const contextManager = new ChatContextManager(
+      savedAgentProfile.profile.system_prompt,
+      sessionMessages,
+      sessionId,
+      ownerData.uuid,
+      sessionCheckpoint,
+      llmUrl,
+      savedAgentProfile.profile.model || DEFAULT_LLM_MODEL,
+      ownerApiKey,
+      new DBCheckpointWriter(db, sessionId),
+      fileManager
+    );
+    if (sessionData.workspace) {
+      const ws = sessionData.workspace;
+      contextManager.setWorkspace(createUserMessage(ws.message, ws.imageB64));
+    }
+
+    const openSession = new OpenSession(
+      db,
+      {} as Agent, // Placeholder - will be replaced after agent creation
+      sessionData,
+      savedAgentProfile,
+      false /* isPersisted */,
+      sessionParticipants,
+      savedAgentProfile.preferences,
+      {} as SkillManager, // Placeholder - will be replaced after agent creation
+      contextManager,
+      connectionManager
+    );
+
+    // Initialize an empty agent (to ensure there are no callbacks before
+    // the OpenSession and client are fully set up).
+
+    const xmcpConfig = Configuration.new(ownerApiKey, xmcpUrl, false);
+    const [agent, skillManager] = await createAgentWithoutSkills(
+      llmUrl,
+      savedAgentProfile.profile,
+      openSession,
+      openSession,
+      contextManager,
+      ownerApiKey,
+      xmcpConfig,
+      undefined,
+      true
+    );
+    await addDefaultChatTools(agent, ownerData.timezone, openSession);
+
+    // Update OpenSession with real agent and skillManager
+    openSession.agent = agent;
+    openSession.skillManager = skillManager;
+
+    await openSession.restoreMcpSettings(
+      savedAgentProfile.profile.mcp_settings
+    );
+    return openSession;
+  }
+
+  static async initWithEmptySession(
+    db: Database,
+    sessionData: SessionData,
+    llmUrl: string,
+    xmcpUrl: string,
+    connectionManager: IUserConnectionManager<ServerToClient>
+  ): Promise<OpenSession> {
+    const sessionMessages: SessionMessage[] = [];
+    const sessionCheckpoint = undefined;
+
+    const { savedAgentProfile, ownerData, ownerApiKey } =
+      await loadDataForEmptySession(
+        db,
+        sessionData.agent_profile_uuid,
+        sessionData.user_uuid
+      );
+    if (!ownerApiKey) {
+      throw new ChatErrorMessage("failed finding owners default api key");
+    }
+
+    const sessionParticipants = new Map<string, TeamParticipant>();
+    sessionParticipants.set(sessionData.user_uuid, {
+      user_uuid: sessionData.user_uuid,
+      nickname: ownerData.nickname || "",
+      email: ownerData.email,
+      role: "owner",
+    });
+
+    return OpenSession.init(
+      db,
+      sessionData,
+      savedAgentProfile,
+      sessionMessages,
+      sessionParticipants,
+      ownerData,
+      ownerApiKey,
+      sessionCheckpoint,
+      llmUrl,
+      xmcpUrl,
+      connectionManager
+    );
+  }
+
+  static async initWithExistingSession(
+    db: Database,
+    sessionId: string,
+    llmUrl: string,
+    xmcpUrl: string,
+    connectionManager: IUserConnectionManager<ServerToClient>
+  ): Promise<OpenSession> {
+    // Load session data from database
+    const {
+      sessionData,
+      sessionMessages,
+      sessionCheckpoint,
+      savedAgentProfile,
+      sessionParticipants,
+      ownerData,
+      ownerApiKey,
+    } = await loadSessionData(db, sessionId);
+    return OpenSession.init(
+      db,
+      sessionData,
+      savedAgentProfile,
+      sessionMessages,
+      sessionParticipants,
+      ownerData,
+      ownerApiKey,
+      sessionCheckpoint,
+      llmUrl,
+      xmcpUrl,
+      connectionManager
+    );
+  }
+
+  onEmpty(): void {
+    logger.info(`[OpenSession.onEmpty] session ${this.sessionUUID}`);
+  }
+
+  getParticipants(): SessionParticipantMap {
+    return this.sessionParticipants;
+  }
+
+  sendSessionData(connectionId: string, clientMessageId: string): void {
+    logger.info(
+      `[SessionRegistry] sending session data for session ${this.sessionUUID}`
+    );
+
+    const sessionInfo = this.serverSessionInfo(clientMessageId);
+    this.connectionManager.sendToConnection(connectionId, sessionInfo);
+
+    // send conversation history
+    const conversationMessages = this.contextManager.getConversationMessages();
+    conversationMessages.forEach((message) => {
+      this.connectionManager.sendToConnection(connectionId, message);
+    });
+
+    // add MCP settings
+    this.sendMcpSettings(connectionId);
+
+    // add system prompt and model
+    const agentProfile = this.agent.getAgentProfile();
+    this.connectionManager.sendToConnection(connectionId, {
+      type: "system_prompt_updated",
+      system_prompt: agentProfile.system_prompt,
+      session_id: this.sessionUUID,
+    });
+    this.connectionManager.sendToConnection(connectionId, {
+      type: "model_updated",
+      model: agentProfile.model || "",
+      session_id: this.sessionUUID,
+    });
+  }
+
+  /**
+   * Restore MCP settings from the database. We only broadcast when
+   * error occurs. The sending of MCP settings is handled by the
+   * the callers of `getMcpSettings` method.
+   * @param mcpServerSettings - MCP server settings from the database.
+   */
+  async restoreMcpSettings(
+    mcpServerSettings: McpServerSettings
+  ): Promise<void> {
+    // TODO: SkillManager.restoreSessings could support a callback to handle
+    // the errors, then we don't need to reimplement the logic.
+
+    // Add each server.  We handle any errors so that conversations may
+    // continue even if not all MCP servers can be enabled (e.g. if the owner
+    // is not present to refresh an oauth token).
+    logger.info(
+      `[OpenSession] restoring MCP settings: ` +
+        JSON.stringify(mcpServerSettings)
+    );
+    const skillManager = this.skillManager;
+    for (const server_name in mcpServerSettings) {
+      try {
+        // TODO: the "[] => all-tools" logic is repeated here, in
+        // `sendSessionData` and in McpServerSettings.  SkillManager could
+        // expose a method `expandMcpServerSettings` which removes missing
+        // elements and expands
+
+        // Handle the case where the mcp server no longer exists
+        if (!skillManager.hasServerBrief(server_name)) {
+          throw new Error(`unknown mcp server: ${server_name}`);
+        }
+
+        // Handle `[]` meaning "all tools"
+        let enabled_tools = mcpServerSettings[server_name];
+        if (enabled_tools.length === 0) {
+          const allTools = await skillManager.getServerTools(server_name);
+          enabled_tools = allTools.map((t) => t.name);
+        }
+
+        await skillManager.addMcpServer(server_name, false);
+        for (const enabled_tool of enabled_tools) {
+          skillManager.enableTool(server_name, enabled_tool);
+        }
+      } catch (e) {
+        this.broadcast({
+          type: "session_error",
+          message: `Error adding MCP server ${server_name}: ${String(e)}`,
+          session_id: this.sessionUUID,
+        });
+        logger.error(`Error adding MCP server ${server_name}: ${String(e)}`);
+      }
+    }
+  }
+
+  /**
+   * Intended to be called within specific `catch` blocks, defining a
+   * consistent error handling approach.  If the exception can be handled it
+   * is and `true` is returned.  Otherwise false is returned (and the caller
+   * is expected to re-throw if they are unable to handle it).
+   */
+  private handleError(err: unknown, from?: string): boolean {
+    const sendError = (msg: ServerSessionError) => {
+      if (from) {
+        this.sendTo(from, msg);
+      } else {
+        this.broadcast(msg);
+      }
+    };
+
+    if (err instanceof ChatFatalError) {
+      // ChatFatalError in session context should send error to specific user
+      // or broadcast if no specific user context
+      sendError({
+        type: "session_error",
+        message: err.message,
+        session_id: this.sessionUUID,
+      });
+    } else if (err instanceof ChatErrorMessage) {
+      sendError({
+        type: "session_error",
+        message: err.message,
+        session_id: this.sessionUUID,
+      });
+    } else {
+      const errString = getErrorString(err);
+      sendError({
+        type: "session_error",
+        message: errString,
+        session_id: this.sessionUUID,
+      });
+    }
+
+    return true;
+  }
+
+  private broadcast(msg: ServerToClient): void {
+    const users: Set<string> = new Set(this.sessionParticipants.keys());
+    this.connectionManager.sendToUsers(users, msg);
+  }
+
+  sendTo(user_uuid: string, msg: ServerToClient): void {
+    this.connectionManager.sendToUsers(new Set([user_uuid]), msg);
+  }
+
+  // IPlatform.openUrl
+  openUrl(url: string, authResultP: Promise<boolean>, display_name: string) {
+    // These requests are always passed to the original owner, since it is
+    // their settings that will be used for all MCP servers.
+    this.broadcast({
+      type: "authentication_started",
+      session_id: this.sessionUUID,
+      url,
+    });
+    this.sendTo(this.userUUID, {
+      type: "authenticate",
+      session_id: this.sessionUUID,
+      url,
+      display_name,
+    });
+
+    // TODO: auth timeout
+    // Don't stall this function waiting for authentication
+    void authResultP.then((result) => {
+      this.sendTo(this.userUUID, {
+        type: "authentication_finished",
+        session_id: this.sessionUUID,
+        url,
+        result,
+      });
+    });
+  }
+
+  // IPlatform.load
+  load(filename: string): Promise<string> {
+    if (process.env.DEVELOPMENT === "1") {
+      return NODE_PLATFORM.load(filename);
+    }
+    throw new ChatErrorMessage("Platform.load not implemented");
+  }
+
+  // IPlatform.renderHTML
+  renderHTML(html: string): Promise<void> {
+    return new Promise<void>((r) => {
+      this.broadcast({
+        type: "render_html",
+        html,
+        session_id: this.sessionUUID,
+      });
+      r();
+    });
+  }
+
+  // IAgentEventHandler.onCompletion
+  onCompletion(result: ChatCompletionAssistantMessageParam): void {
+    logger.debug(`[OpenSession.onCompletion] : ${JSON.stringify(result)}`);
+    // Nothing to broadcast. Caller will receive this via onAgentMessage.
+    this.contextManager.processAgentResponse(result);
+  }
+
+  // IAgentEventHandler.onToolCallResult
+  onToolCallResult(result: ChatCompletionToolMessageParam): void {
+    logger.debug(`[onToolCallResult] : ${JSON.stringify(result)}`);
+    const toolCallMessage = this.contextManager.processToolCallResult(result);
+    this.broadcast(toolCallMessage);
+  }
+
+  // IAgentEventHandler.onToolCall
+  async onToolCall(
+    toolCall: ChatCompletionMessageToolCall,
+    agentTool: boolean
+  ): Promise<boolean> {
+    if (agentTool) {
+      // "Agent" tools are considered internal to the agent, and are always
+      // permitted.  Inform all clients and immediately approve.
+      this.broadcast({
+        type: "tool_call",
+        tool_call: toolCall,
+        session_id: this.sessionUUID,
+      });
+      return true;
+    }
+
+    // TODO: Need a proper mapping to/from MCP calls to tool names
+
+    const [serverName, tool] = toolCall.function.name.split("__");
+    const autoApproved = prefsGetAutoApprove(
+      this.agentProfilePreferences,
+      serverName,
+      tool
+    );
+    if (!autoApproved) {
+      const { id, resultP } = this.approvalManager.startApproval(
+        toolCall.function.name
+      );
+      this.broadcast({
+        type: "approve_tool_call",
+        id,
+        tool_call: toolCall,
+        session_id: this.sessionUUID,
+      });
+
+      try {
+        logger.debug(`[OpenSession.onToolCall] awaiting approval ${id}`);
+        const { approved, auto_approve } = await resultP;
+        logger.debug(
+          `[OpenSession.onToolCall] approval ${id}: ${String(approved)}`
+        );
+        if (auto_approve) {
+          logger.debug(
+            "[OpenSession.onToolCall] auto_approve set. updated preferences"
+          );
+          const autoApprovalMsg = await this.onSetAutoApproval(
+            serverName,
+            tool,
+            true
+          );
+          if (autoApprovalMsg) {
+            this.broadcast(autoApprovalMsg);
+          }
+        }
+
+        return approved;
+      } catch (e) {
+        logger.debug(
+          `[OpenSession.onToolCall] error waiting for approval ${id}: ` +
+            String(e)
+        );
+        return false;
+      }
+    } else {
+      this.broadcast({
+        type: "tool_call",
+        tool_call: toolCall,
+        session_id: this.sessionUUID,
+      });
+      return true;
+    }
+  }
+
+  // IAgentEventHandler.onAgentMessage
+  // eslint-disable-next-line @typescript-eslint/require-await
+  async onAgentMessage(msg: string, end: boolean): Promise<void> {
+    logger.debug(
+      `[OpenSession.onAgentMessage] msg: ${msg}, end: ${String(end)}`
+    );
+
+    // Inform the contextManager and broadcast the ServerAgentMessageChunk
+    const agentMsgChunk = this.contextManager.processAgentMessage(msg, end);
+    this.broadcast(agentMsgChunk);
+  }
+
+  /**
+   * Handle incoming session-related message from client.
+   * These messages include tool calls, MCP changes, etc.
+   * These messages are processed by `messageQueue`.
+   */
+  onClientSessionMessage(from: string, message: ClientSessionMessage): void {
+    logger.info(
+      `[OpenSession] Message from ${from} in session ` +
+        `${this.sessionUUID}: ${JSON.stringify(message)}`
+    );
+
+    // Validate user is in session
+    if (!this.sessionParticipants.get(from)) {
+      logger.warn(
+        `User ${from} not in session ${this.sessionUUID} - ignoring message`
+      );
+      this.sendTo(from, {
+        type: "session_error",
+        message: "You are not a participant in this session",
+        session_id: this.sessionUUID,
+        client_message_id: message.client_message_id,
+      });
+      return;
+    }
+
+    // Enqueue message for processing
+    if (!this.messageQueue.tryEnqueue({ msg: message, from })) {
+      this.sendTo(
+        from,
+        this.addSessionContext({
+          type: "session_error",
+          message: "message queue full. try again later",
+          client_message_id: message.client_message_id,
+        })
+      );
+    }
+  }
+
+  /**
+   * Get MCP settings in the form of `ServerToClient` messages.
+   * This is used to restore MCP settings when a user joins a session.
+   */
+  private sendMcpSettings(connectionId: string): void {
+    const servers = this.skillManager.getMcpServerNames();
+    servers.forEach((server_name) => {
+      const mcpServer = this.skillManager.getMcpServer(server_name);
+      const tools = mcpServer.getTools();
+      const enabled_tools = Array.from(mcpServer.getEnabledTools().keys());
+      this.connectionManager.sendToConnection(connectionId, {
+        type: "mcp_server_added",
+        server_name,
+        tools,
+        enabled_tools,
+        session_id: this.sessionUUID,
+      });
+    });
+  }
+
+  /**
+   * Called once for each message.
+   */
+  private async processMessage(
+    queuedMessage: QueuedClientMessage
+  ): Promise<void> {
+    logger.debug(
+      `[onMessage]: processing (${queuedMessage.from}) ` +
+        JSON.stringify(queuedMessage.msg)
+    );
+
+    // 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 = this.onUserMessage(msg, 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 "tool_call_approval_result":
+          if (
+            this.approvalManager.approvalResult(
+              msg.id,
+              msg.result,
+              msg.auto_approve
+            )
+          ) {
+            broadcastMsg = {
+              type: "tool_call_approval_result",
+              id: msg.id,
+              result: msg.result,
+              session_id: this.sessionUUID,
+            };
+          }
+          break;
+        case "set_auto_approval":
+          broadcastMsg = await this.onSetAutoApproval(
+            msg.server_name,
+            msg.tool,
+            msg.auto_approve
+          );
+          break;
+        case "set_system_prompt":
+          broadcastMsg = await this.onSetSystemPrompt(msg.system_prompt);
+          break;
+        case "set_model":
+          broadcastMsg = await this.onSetModel(msg.model);
+          break;
+        case "set_workspace":
+          await this.onSetWorkspace(msg, queuedMessage.from);
+          break;
+        case "set_mcp_server_config":
+          await this.handleSetMcpServerConfig(msg, queuedMessage.from);
+          break;
+        default: {
+          const exhaustive: never = msg; // Error => non-exhaustive switch-case.
+          return exhaustive;
+        }
+      }
+
+      if (broadcastMsg) {
+        if (broadcastMsg instanceof Array) {
+          broadcastMsg.map((msg) => {
+            this.broadcast(msg);
+          });
+        } else {
+          this.broadcast(broadcastMsg);
+        }
+      }
+    } catch (err: unknown) {
+      if (!this.handleError(err, queuedMessage.from)) {
+        throw err;
+      }
+    }
+  }
+
+  private async onSetWorkspace(
+    msg: ClientSetWorkspace,
+    sender: string
+  ): Promise<void> {
+    this.contextManager.setWorkspace(
+      createUserMessage(msg.message, msg.imageB64, sender)
+    );
+    const m = this.workspaceUserMessageData();
+
+    // It's possible that the session has not been written to the DB yet (if
+    // the workspace is updated before any messages are sent).
+
+    if (!this.isPersisted) {
+      await this.createSessionInDB();
+    } else {
+      await this.db.sessionUpdateWorkspace(this.sessionUUID, m);
+    }
+    assert(this.isPersisted);
+  }
+
+  private async processUserMessages(msgs: ServerUserMessage[]): Promise<void> {
+    if (msgs.length === 0) {
+      logger.debug("ignoring empty messages");
+      return;
+    }
+
+    // Begin an agent response using the accumulated user messages.
+
+    const { llmUserMessages, agentFirstChunk } =
+      this.contextManager.startAgentResponse(msgs);
+    this.broadcast(agentFirstChunk);
+
+    try {
+      await this.agent.userMessagesRaw(llmUserMessages);
+    } catch (e) {
+      logger.warn(
+        `[OpenSession.processUserMessages] agent error: ${String(e)}`
+      );
+
+      // Errors during agent replies must be turned into messages.
+
+      const errMsg = `error from LLM: ${String(e)}`;
+      await this.onAgentMessage(errMsg, true);
+      const err = this.contextManager.revertAgentResponse(errMsg);
+      this.broadcast(err);
+      return;
+    }
+
+    const newSessionMessages = this.contextManager.endAgentResponse();
+
+    logger.debug(
+      "[processUserMessages] newSessionMessages: " +
+        JSON.stringify(newSessionMessages)
+    );
+
+    try {
+      // Append to in-memory conversation and write to the DB
+      await this.db.sessionMessagesAppend(this.sessionUUID, newSessionMessages);
+    } catch (e) {
+      if (!this.handleError(e)) {
+        throw e;
+      }
+    }
+  }
+
+  private onUserMessage(
+    msg: ClientUserMessage,
+    from: string
+  ): ServerUserMessage | undefined {
+    // Return a ServerUserMessage for broadcast.  The actual message is places
+    // on a queue to be dealt with in another loop.  This allows Agent
+    // processing of user messages to depend on other messages.
+
+    assert(msg);
+    assert(from);
+
+    // Assign the user message_idx and attempt to enqueue.
+
+    const userMessage = this.contextManager.processUserMessage(msg, from);
+
+    // Special case for the first message of the session
+
+    if (userMessage.message_idx === MESSAGE_INDEX_START_VALUE) {
+      // No need to wait for this to complete before broadcasting.
+      void this.onFirstMessage(userMessage);
+    }
+
+    if (!this.userMessageQueue.tryEnqueue(userMessage)) {
+      // We failed to enqueue - revert the `getNextMessageIdx`
+
+      // NOTE: Nothing should await between `getNextMessageIdx` and
+      // `freeMessageIdx` here.
+
+      this.contextManager.unprocessUserMessage(userMessage);
+
+      this.sendTo(from, {
+        type: "session_error",
+        message: "failed to queue message.  try again later.",
+        session_id: this.sessionUUID,
+      });
+      return undefined;
+    }
+
+    // Message is enqueued.  Broadcast to all participants.
+    return userMessage;
+  }
+
+  private async createSessionInDB(): Promise<void> {
+    assert(!this.isPersisted);
+    const sessionCreateData: SessionCreateData = {
+      session_uuid: this.sessionUUID,
+      title: this.sessionTitle,
+      team_uuid: this.teamUUID,
+      agent_profile_uuid: this.agentProfileUUID,
+      workspace: this.workspaceUserMessageData(),
+      user_uuid: this.userUUID,
+    };
+    logger.info(
+      `[OpenSession.onFirstMessage] writing session ${this.sessionUUID}`
+    );
+
+    await this.db.sessionCreate(sessionCreateData);
+    this.isPersisted = true;
+  }
+
+  private async onFirstMessage(userMsg: ServerUserMessage): Promise<void> {
+    // Update title on the class before writing to the DB
+
+    this.sessionTitle = userMsg.message?.slice(0, 128) || "New Chat";
+
+    // The session may already have been saved (e.g. if the workspace is
+    // updated before any messages are sent).
+
+    if (!this.isPersisted) {
+      await this.createSessionInDB();
+    } else {
+      await this.db.sessionUpdateTitle(this.sessionUUID, this.sessionTitle);
+    }
+    assert(this.isPersisted);
+
+    // Broadcast the SessionUpdated message
+
+    const msg: ServerSessionUpdate = {
+      type: "session_update",
+      session_id: this.sessionUUID,
+      title: this.sessionTitle,
+    };
+    this.broadcast(msg);
+  }
+
+  private async onAddMcpServer(
+    serverName: string,
+    enableAll: boolean
+  ): Promise<ServerMcpServerAdded> {
+    logger.info(
+      `[onAddMcpServer]: Adding server ${serverName} ` +
+        `(enable_all: ${String(enableAll)})`
+    );
+
+    if (this.skillManager.hasMcpServer(serverName)) {
+      throw new ChatErrorMessage(`${serverName} already added`);
+    }
+
+    if (!this.skillManager.hasServerBrief(serverName)) {
+      throw new ChatErrorMessage(`no such server: ${serverName}`);
+    }
+
+    await this.skillManager.addMcpServer(serverName, enableAll);
+    this.skillManager.enableAllTools(serverName);
+
+    // Save changes to the AgentProfile
+
+    await this.updateAgentProfile();
+
+    // Broadcast the message to all participants.
+
+    const server = this.skillManager.getMcpServer(serverName);
+    const tools = server.getTools();
+    const enabled_tools = Array.from(server.getEnabledTools().keys());
+    return {
+      type: "mcp_server_added",
+      server_name: serverName,
+      tools,
+      enabled_tools,
+      session_id: this.sessionUUID,
+    };
+  }
+
+  private async onRemoveMcpServer(
+    server_name: string
+  ): Promise<ServerMcpServerRemoved> {
+    logger.info(`[onRemoveMcpServer]: Removing server ${server_name}`);
+
+    if (!this.skillManager.hasMcpServer(server_name)) {
+      throw new ChatErrorMessage(`${server_name} not enabled`);
+    }
+
+    await this.skillManager.removeMcpServer(server_name);
+    await this.updateAgentProfile();
+
+    return {
+      type: "mcp_server_removed",
+      server_name,
+      session_id: this.sessionUUID,
+    };
+  }
+
+  private async onEnableMcpServerTool(
+    server_name: string,
+    tool: string
+  ): Promise<ServerMcpServerToolEnabled> {
+    this.ensureMcpServer(server_name);
+    this.skillManager.enableTool(server_name, tool);
+
+    await this.updateAgentProfile();
+
+    return {
+      type: "mcp_server_tool_enabled",
+      server_name,
+      tool,
+      session_id: this.sessionUUID,
+    };
+  }
+
+  private async onDisableMcpServerTool(
+    server_name: string,
+    tool: string
+  ): Promise<ServerMcpServerToolDisabled> {
+    this.ensureMcpServerAndTool(server_name, tool);
+    this.skillManager.disableTool(server_name, tool);
+
+    await this.updateAgentProfile();
+
+    return {
+      type: "mcp_server_tool_disabled",
+      server_name,
+      tool,
+      session_id: this.sessionUUID,
+    };
+  }
+
+  private 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 server = this.ensureMcpServer(server_name);
+    const enabledTools = server.getEnabledTools();
+    const msgs: ServerMcpServerToolEnabled[] = [];
+    for (const tool of server.getTools()) {
+      if (!enabledTools.get(tool.name)) {
+        this.skillManager.enableTool(server_name, tool.name);
+        msgs.push({
+          type: "mcp_server_tool_enabled",
+          server_name,
+          tool: tool.name,
+          session_id: this.sessionUUID,
+        });
+      }
+    }
+
+    await this.updateAgentProfile();
+
+    return msgs;
+  }
+
+  private 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 server = this.ensureMcpServer(server_name);
+    const enabledTools = server.getEnabledTools();
+    const msgs: ServerMcpServerToolDisabled[] = [];
+    for (const tool of enabledTools.keys()) {
+      this.skillManager.disableTool(server_name, tool);
+      msgs.push({
+        type: "mcp_server_tool_disabled",
+        server_name,
+        tool,
+        session_id: this.sessionUUID,
+      });
+    }
+
+    await this.updateAgentProfile();
+
+    return msgs;
+  }
+
+  private async onSetSystemPrompt(
+    system_prompt: string
+  ): Promise<ServerSystemPromptUpdated> {
+    this.agent.setSystemPrompt(system_prompt);
+    await this.updateAgentProfile();
+    return {
+      type: "system_prompt_updated",
+      system_prompt,
+      session_id: this.sessionUUID,
+    };
+  }
+
+  private async onSetModel(model: string): Promise<ServerModelUpdated> {
+    this.agent.setModel(model);
+    await this.updateAgentProfile();
+    return { type: "model_updated", model, session_id: this.sessionUUID };
+  }
+
+  private async handleSetMcpServerConfig(
+    msg: ClientSetMcpServerConfig,
+    from: string
+  ): Promise<void> {
+    if (!this.skillManager.hasServerBrief(msg.server_name)) {
+      throw new Error(`unknown mcp server: ${msg.server_name}`);
+    }
+
+    const msc = this.db.createTypedClient(DbMcpServerConfigs);
+    if (msg.config) {
+      await msc.setConfigForUser(from, msg.server_name, msg.config);
+    } else {
+      await msc.deleteConfigForUser(from, msg.server_name);
+    }
+
+    // If the mcp server was enabled, disable and re-enable to ensure the new
+    // settings are reflected.  Note, we must track the set of enabled tools.
+
+    if (this.skillManager.hasMcpServer(msg.server_name)) {
+      const enabledTools = this.skillManager
+        .getMcpServer(msg.server_name)
+        .getEnabledTools();
+      await this.skillManager.removeMcpServer(msg.server_name);
+      await this.skillManager.addMcpServer(msg.server_name, false);
+      for (const tool of enabledTools.keys()) {
+        this.skillManager.enableTool(msg.server_name, tool);
+      }
+    }
+
+    // TODO: Do we want to braodcast an "mcp_server_config_updated" message?
+  }
+
+  private async onSetAutoApproval(
+    serverName: string,
+    tool: string,
+    autoApprove: boolean
+  ): Promise<ServerToolAutoApprovalSet | undefined> {
+    if (
+      prefsSetAutoApprove(
+        this.agentProfilePreferences,
+        serverName,
+        tool,
+        autoApprove
+      )
+    ) {
+      await this.db.updateAgentProfilePreferences(
+        this.agentProfileUUID,
+        this.agentProfilePreferences
+      );
+
+      return {
+        type: "tool_auto_approval_set",
+        server_name: serverName,
+        tool,
+        auto_approve: autoApprove,
+        session_id: this.sessionUUID,
+      };
+    }
+  }
+
+  private ensureMcpServer(serverName: string): McpServerInfo {
+    return this.skillManager.getMcpServer(serverName);
+  }
+
+  private ensureMcpServerAndTool(serverName: string, toolName: string): Tool {
+    const server = this.ensureMcpServer(serverName);
+    const tool = server.getTool(toolName);
+    if (!tool) {
+      throw new ChatErrorMessage(`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);
+  }
+
+  /**
+   * Update participant in session (called by SessionRegistry).
+   * This only updates the local participant map - actual membership
+   * tracking is handled by SessionRegistry.
+   */
+  addParticipant(userId: string, role: TeamParticipant): void {
+    this.sessionParticipants.set(userId, role);
+    // Broadcast result to all session participants
+    const broadcastMessage: ServerUserAdded = {
+      type: "user_added",
+      user_uuid: userId,
+      role: "participant",
+      nickname: role.nickname,
+      email: role.email,
+      session_id: this.sessionUUID,
+    };
+
+    this.broadcast(broadcastMessage);
+  }
+
+  /**
+   * Remove participant from session (called by SessionRegistry).
+   * This only updates the local participant map - actual membership
+   * tracking is handled by SessionRegistry.
+   */
+  removeParticipant(userId: string): void {
+    this.sessionParticipants.delete(userId);
+    // Broadcast result to all session participants
+    const broadcastMessage: ServerUserRemoved = {
+      type: "user_removed",
+      user_uuid: userId,
+      session_id: this.sessionUUID,
+    };
+
+    this.broadcast(broadcastMessage);
+  }
+
+  private getSessionParticipants(): TeamParticipant[] {
+    return Array.from(this.sessionParticipants.entries()).map(
+      ([userId, user]) => ({
+        nickname: user.nickname,
+        email: user.email,
+        user_uuid: userId,
+        role: user.role,
+      })
+    );
+  }
+
+  private workspaceUserMessageData(): UserMessageData | undefined {
+    const workspaceMsg = this.contextManager.getWorkspace();
+    return workspaceMsg
+      ? llmUserMessageToUserMessageData(workspaceMsg)
+      : undefined;
+  }
+
+  serverSessionInfo(clientMessageId: string): ServerSessionInfo {
+    return {
+      type: "session_info",
+      owner_uuid: this.userUUID,
+      title: this.sessionTitle,
+      saved_agent_profile: this.savedAgentProfile,
+      workspace: this.workspaceUserMessageData(),
+      updated_at: this.sessionUpdatedAt,
+      participants: this.getSessionParticipants(),
+      mcp_server_briefs: this.skillManager.getServerBriefs(),
+      agent_preferences: this.agentProfilePreferences,
+      client_message_id: clientMessageId,
+      session_id: this.sessionUUID,
+      team_uuid: this.teamUUID,
+    };
+  }
+
+  // Helper method to add session context to messages
+  private addSessionContext<T extends object>(
+    msg: T
+  ): T & { session_id: string } {
+    return {
+      ...msg,
+      session_id: this.sessionUUID,
+    };
+  }
+}
+
+async function loadDataForEmptySession(
+  db: Database,
+  agent_profile_uuid: string,
+  owner_uuid: string
+): Promise<{
+  savedAgentProfile: SavedAgentProfile;
+  ownerData: UserData;
+  ownerApiKey: string;
+}> {
+  const [savedAgentProfile, ownerData, ownerApiKey] = await Promise.all([
+    db.getSavedAgentProfileById(agent_profile_uuid),
+    db.getUserFromUuid(owner_uuid),
+    db.getUserApiKey(owner_uuid),
+  ]);
+  if (!savedAgentProfile) {
+    throw new ChatFatalError(`No such agent profile: ${agent_profile_uuid}`);
+  }
+  if (!ownerData) {
+    throw new ChatFatalError(`No owner ${owner_uuid}`);
+  }
+  if (!ownerData.nickname) {
+    throw new ChatFatalError(
+      `User ${owner_uuid} has no username - cannot create session`
+    );
+  }
+  if (!ownerApiKey) {
+    throw new ChatFatalError(
+      `User ${owner_uuid} has no api keys - cannot create session`
+    );
+  }
+
+  return { savedAgentProfile, ownerData, ownerApiKey };
+}
+
+/**
+ * Loads session data, agent profile, and owner information from the database.
+ */
+async function loadSessionData(
+  db: Database,
+  sessionId: string
+): Promise<{
+  sessionData: SessionData;
+  sessionMessages: SessionMessage[];
+  sessionCheckpoint: SessionCheckpoint | undefined;
+  savedAgentProfile: SavedAgentProfile;
+  ownerData: UserData;
+  ownerApiKey: string;
+  sessionParticipants: SessionParticipantMap;
+}> {
+  const [sessionData, sessionMessages, sessionParticipants, sessionCheckpoint] =
+    await Promise.all([
+      db.sessionGetById(sessionId),
+      db.sessionMessagesGetConversation(sessionId, DEFAULT_NUM_MESSGAES, 0),
+      db.sessionGetParticipants(sessionId),
+      db.sessionCheckpointGet(sessionId),
+    ]);
+  if (!sessionData) {
+    throw new ChatFatalError(`No such session: ${sessionId}`);
+  }
+
+  const { savedAgentProfile, ownerData, ownerApiKey } =
+    await loadDataForEmptySession(
+      db,
+      sessionData.agent_profile_uuid,
+      sessionData.user_uuid
+    );
+
+  return {
+    sessionData,
+    sessionMessages,
+    sessionCheckpoint,
+    savedAgentProfile,
+    ownerData,
+    ownerApiKey,
+    sessionParticipants: createSessionParticipantMap(sessionParticipants),
+  };
+}