@xalia/agent 0.5.8 → 0.6.0

package/src/chat/client/chatClient.ts

@@ -0,0 +1,900 @@
+import { v4 as uuidv4 } from "uuid";
+
+import { getLogger, SavedAgentProfile } from "@xalia/xmcp/sdk";
+
+import { Connection } from "./connection";
+import { SessionClient } from "./sessionClient";
+import { IChatClientEventHandler } from "./interfaces";
+import {
+  SessionData,
+  TeamInfo,
+  TeamParticipant,
+  AgentSessionData,
+} from "../data/dataModels";
+import { createSessionParticipantMap } from "../data/database";
+import {
+  ServerSessionScopedMessage,
+  ServerSessionInfo,
+  isServerSessionScopedMessage,
+  ServerControlError,
+  isServerControlMessage,
+  ServerControlMessage,
+  ServerControlSessionDeleted,
+  ServerControlTeamCreated,
+  ServerToClient,
+  ClientToServer,
+  ServerControlAgentProfileCreated,
+  ServerControlAgentProfileDeleted,
+} from "../protocol/messages";
+import { SESSION_CREATE_TIMEOUT, SESSION_JOIN_TIMEOUT } from "./constants";
+import { ITeamManager } from "./teamManager";
+import { buildAgentSessionMap } from "../utils/agentSessionMap";
+
+const logger = getLogger();
+
+// client team info, used for session tracking etc.
+export type ClientTeamInfo = {
+  team_uuid: string;
+  team_name: string;
+  owner_uuid: string;
+  participants: Map<string, TeamParticipant>;
+  agentSessionMap: Map<string, AgentSessionData>;
+};
+
+type SessionJoinOrCreateResolver = {
+  resolve: (
+    sessionClient: SessionClient,
+    sessionInfo: ServerSessionInfo
+  ) => void;
+  reject: (error: Error) => void;
+  sessionId: string;
+  clientMessageId: string;
+};
+
+export class ChatClient implements ITeamManager {
+  private constructor(
+    private connection: Connection<ClientToServer, ServerToClient>,
+    private eventHandler: IChatClientEventHandler,
+    // agent_uuid -> agent session data
+    private userAgentSessionMap: Map<string, AgentSessionData> = new Map(),
+    // team_uuid -> team info
+    private teams: Map<string, ClientTeamInfo> = new Map(),
+    private closed: boolean = false,
+    private currentSessionId: string | undefined = undefined,
+    // note: currentTeamId will not be reset when swtiching to a user session.
+    private currentTeamId: string | undefined = undefined,
+    // session_uuid -> session client
+    private activeSessions: Map<string, SessionClient> = new Map(),
+    // session_uuid -> true if there is a new message.
+    private newMessage: Map<string, boolean> = new Map(),
+    // Track pending session join request, this should be a singleton
+    private sessionJoinOrCreateRes:
+      | SessionJoinOrCreateResolver
+      | undefined = undefined
+  ) {}
+
+  /**
+   * Connect to server and create ChatClient instance
+   */
+  static async init(
+    url: string,
+    token: string,
+    eventHandler: IChatClientEventHandler
+  ): Promise<ChatClient> {
+    const connection = new Connection<ClientToServer, ServerToClient>({
+      url,
+      token,
+    });
+    return new Promise((resolveClient, reject) => {
+      let client: ChatClient | undefined = undefined;
+      const clientMessageId = uuidv4();
+
+      // Register session_info handler for initialization
+      connection.on("control_session_list", (msg) => {
+        // get user sessions, user agents, and team sessions
+        const userSessions = new Map<string, SessionData>();
+        msg.user_sessions.forEach((session) => {
+          userSessions.set(session.session_uuid, session);
+        });
+
+        const userAgents = new Map<string, SavedAgentProfile>();
+        msg.user_agents.forEach((agent) => {
+          userAgents.set(agent.uuid, agent);
+        });
+
+        const userAgentSessionMap = buildAgentSessionMap(
+          userSessions,
+          userAgents
+        );
+        logger.info(
+          `[ChatClient.init] session list msg: ${JSON.stringify(msg)}`
+        );
+        for (const agentSession of userAgentSessionMap.values()) {
+          logger.info(
+            `[ChatClient.init] agent: ` +
+              agentSession.agent_profile.profile_name
+          );
+        }
+
+        const teams = new Map<string, ClientTeamInfo>();
+        msg.team_sessions.forEach((team) => {
+          teams.set(team.team_uuid, ChatClient.toClientTeamInfo(team));
+        });
+
+        if (!client) {
+          logger.info("Creating ChatClient");
+          client = new ChatClient(
+            connection,
+            eventHandler,
+            userAgentSessionMap,
+            teams
+          );
+          resolveClient(client);
+        } else {
+          client.userAgentSessionMap = userAgentSessionMap;
+          client.teams = teams;
+          void client.eventHandler.onMessage(msg, client);
+        }
+      });
+
+      // Register general message handler for established client
+      connection.on("*", (msg) => {
+        if (!client) {
+          if (msg.type === "control_error") {
+            logger.error(
+              `[ChatClient] client error (init): ${JSON.stringify(msg)}`
+            );
+            reject(new Error(msg.message));
+          } else {
+            logger.error(
+              `[ChatClient] invalid message (init): ${JSON.stringify(msg.type)}`
+            );
+          }
+          return;
+        }
+        if (isServerControlMessage(msg)) {
+          client.handleControlMessage(msg);
+        } else if (isServerSessionScopedMessage(msg)) {
+          client.handleSessionMessage(msg);
+        } else {
+          logger.error(`[ChatClient] unhandled msg: ${JSON.stringify(msg)}`);
+        }
+      });
+
+      // Register direct error callback for client-side errors.
+      connection.onError((errorMessage) => {
+        if (client) {
+          client.connection.close();
+        }
+        eventHandler.onError(`Connection error: ${errorMessage}`);
+      });
+
+      // Start connection
+      connection
+        .connect()
+        .then(() => {
+          connection.send({
+            type: "control_get_session_list",
+            client_message_id: clientMessageId,
+          });
+        })
+        .catch((error: unknown) => {
+          reject(error instanceof Error ? error : new Error(String(error)));
+        });
+    });
+  }
+
+  private static toClientTeamInfo(teamInfo: TeamInfo): ClientTeamInfo {
+    const sessions = new Map(
+      teamInfo.sessions.map((session) => [session.session_uuid, session])
+    );
+    const agents = new Map(teamInfo.agents.map((agent) => [agent.uuid, agent]));
+    const agentSessionMap = buildAgentSessionMap(sessions, agents);
+    return {
+      team_uuid: teamInfo.team_uuid,
+      team_name: teamInfo.team_name,
+      owner_uuid: teamInfo.owner_uuid,
+      participants: new Map(
+        teamInfo.participants.map((participant) => [
+          participant.user_uuid,
+          participant,
+        ])
+      ),
+      agentSessionMap,
+    };
+  }
+
+  /**
+   * Get the current session.
+   * Will throw an error if the current session is not set.
+   * @param location - The location of the caller.
+   * @returns The current session client.
+   */
+  public getCurrentSession(location: string): SessionClient {
+    const session = this.getCurrentSessionMaybe();
+    if (!session) {
+      throw new Error(`[${location}]: No current session`);
+    }
+    return session;
+  }
+
+  public getCurrentSessionMaybe(): SessionClient | undefined {
+    if (!this.currentSessionId) {
+      return undefined;
+    }
+    return this.activeSessions.get(this.currentSessionId);
+  }
+
+  public getCurrentAgentUuid(): string | undefined {
+    const session = this.getCurrentSessionMaybe();
+    if (!session) {
+      return undefined;
+    }
+    return session.getAgentUuid();
+  }
+
+  getUserAgentSessionMap(): Map<string, AgentSessionData> {
+    return this.userAgentSessionMap;
+  }
+
+  getClientTeamInfo(teamUuid: string): ClientTeamInfo | undefined {
+    return this.teams.get(teamUuid);
+  }
+
+  // ITeamManager -> getCurrentTeamInfo
+  getCurrentTeamInfo(): ClientTeamInfo | undefined {
+    if (!this.currentTeamId) {
+      return undefined;
+    }
+    return this.teams.get(this.currentTeamId);
+  }
+
+  // ITeamManager -> getTeams
+  getTeams(): Map<string, ClientTeamInfo> {
+    return this.teams;
+  }
+
+  /**
+   * Get the currently active session ID
+   */
+  getCurrentSessionId(): string | undefined {
+    return this.currentSessionId;
+  }
+
+  /**
+   * Get the currently team ID
+   */
+  getCurrentTeamId(): string | undefined {
+    return this.currentTeamId;
+  }
+
+  // ITeamManager -> setCurrentTeamId
+  setCurrentTeamId(teamUuid: string): void {
+    if (!this.teams.has(teamUuid)) {
+      throw new Error(`Team ${teamUuid} not found`);
+    }
+    this.currentTeamId = teamUuid;
+  }
+
+  /**
+   * Check if client is connected and operational
+   */
+  isClosed(): boolean {
+    return this.closed;
+  }
+
+  shutdown(): void {
+    this.connection.close();
+    this.closed = true;
+    this.activeSessions.clear();
+    this.newMessage.clear();
+    this.teams.clear();
+    this.currentSessionId = undefined;
+    this.currentTeamId = undefined;
+    this.sessionJoinOrCreateRes = undefined;
+  }
+
+  public async createNewSession(
+    newTitle: string,
+    agentProfileId?: string,
+    teamId?: string
+  ): Promise<SessionClient> {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+
+    if (this.sessionJoinOrCreateRes) {
+      logger.error(
+        `[ChatClient] creating new session with ` +
+          `profile id ${JSON.stringify(agentProfileId)}`
+      );
+      throw new Error("Session join/create is already in progress");
+    }
+
+    return new Promise((resolve, reject) => {
+      const clientMessageId = uuidv4();
+
+      // Set up timeout for the request
+      const timeoutId = setTimeout(() => {
+        this.sessionJoinOrCreateRes = undefined;
+        reject(new Error(`Session join timeout for creating new session`));
+      }, SESSION_CREATE_TIMEOUT);
+
+      // Track this pending request
+      this.sessionJoinOrCreateRes = {
+        resolve: (
+          sessionClient: SessionClient,
+          sessionInfo: ServerSessionInfo
+        ) => {
+          const sessionId = this.sessionJoinOrCreateRes?.sessionId;
+          if (!sessionId) {
+            throw new Error("Session id is not set in resolving.");
+          }
+          if (sessionId !== sessionInfo.session_id) {
+            throw new Error("SessionInfo id mismatch");
+          }
+          clearTimeout(timeoutId);
+          this.activeSessions.set(sessionId, sessionClient);
+          this.currentSessionId = sessionId;
+          if (sessionInfo.team_uuid) {
+            this.currentTeamId = sessionInfo.team_uuid;
+          }
+          this.sessionJoinOrCreateRes = undefined;
+          this.addSessionToAgentSessionMap(sessionInfo);
+          resolve(sessionClient);
+        },
+        reject: (error: Error) => {
+          clearTimeout(timeoutId);
+          this.sessionJoinOrCreateRes = undefined;
+          reject(error);
+        },
+        sessionId: "",
+        clientMessageId,
+      };
+
+      // Send session_create message
+      try {
+        this.connection.send({
+          type: "control_session_create",
+          client_message_id: clientMessageId,
+          title: newTitle,
+          agent_profile_id: agentProfileId,
+          team_id: teamId,
+        });
+        logger.info(
+          `[ChatClient] Sent session_create for session ${newTitle}` +
+            ` client message id: ${clientMessageId}`
+        );
+      } catch (error) {
+        this.sessionJoinOrCreateRes = undefined;
+        clearTimeout(timeoutId);
+        reject(error instanceof Error ? error : new Error(String(error)));
+      }
+    });
+  }
+
+  /**
+   * Connect to an existing session by sending session_join message
+   * and waiting for session_info response from server
+   */
+  async connectToSession(sessionId: string): Promise<SessionClient> {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+
+    if (this.sessionJoinOrCreateRes) {
+      logger.error(`[ChatClient] connecting to session ${sessionId}`);
+      throw new Error("Session join/create is already in progress");
+    }
+
+    // TODO: we can directly return a session client if we cache conversation.
+    // For now, we go through the join process.
+    return new Promise((resolve, reject) => {
+      const clientMessageId = uuidv4();
+
+      // Set up timeout for the request
+      const timeoutId = setTimeout(() => {
+        this.sessionJoinOrCreateRes = undefined;
+        reject(new Error(`Session join timeout for session ${sessionId}`));
+      }, SESSION_JOIN_TIMEOUT);
+
+      // Track this pending request
+      this.sessionJoinOrCreateRes = {
+        resolve: (
+          sessionClient: SessionClient,
+          sessionInfo: ServerSessionInfo
+        ) => {
+          clearTimeout(timeoutId);
+          this.activeSessions.set(sessionId, sessionClient);
+          this.currentSessionId = sessionId;
+          if (sessionInfo.team_uuid) {
+            this.currentTeamId = sessionInfo.team_uuid;
+          }
+          this.sessionJoinOrCreateRes = undefined;
+          this.updateSessionInfo(sessionInfo);
+          logger.info(`[ChatClient] joined session ${sessionId}`);
+          resolve(sessionClient);
+        },
+        reject: (error: Error) => {
+          clearTimeout(timeoutId);
+          this.sessionJoinOrCreateRes = undefined;
+          logger.error(
+            `[ChatClient] failed to join session` +
+              ` ${sessionId}: ${error.message}`
+          );
+          reject(error);
+        },
+        sessionId,
+        clientMessageId,
+      };
+
+      // Send session_join message
+      try {
+        this.connection.send({
+          type: "control_session_join",
+          client_message_id: clientMessageId,
+          target_session_id: sessionId,
+        });
+        logger.debug(`[ChatClient] Sent session_join for session ${sessionId}`);
+      } catch (error) {
+        this.sessionJoinOrCreateRes = undefined;
+        clearTimeout(timeoutId);
+        reject(error instanceof Error ? error : new Error(String(error)));
+      }
+    });
+  }
+
+  /**
+   * Fetch the session list from the server
+   */
+  fetchSessionList(): void {
+    this.connection.send({
+      type: "control_get_session_list",
+      client_message_id: uuidv4(),
+    });
+  }
+
+  /**
+   * Delete a session by sending session_delete_request message
+   * and waiting for session_deleted response from server
+   * @param sessionId
+   * @returns void
+   */
+  deleteSession(sessionId: string): void {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+
+    this.connection.send({
+      type: "control_session_delete",
+      client_message_id: uuidv4(),
+      target_session_id: sessionId,
+    });
+
+    logger.debug(
+      `[ChatClient] Sent session_delete_request for` + ` session ${sessionId}`
+    );
+  }
+
+  /**
+   * Create a new team with initial members
+   * @param teamName
+   * @param initialMembers could be UUIDs or emails
+   */
+  createNewTeam(teamName: string, initialMembers: string[]): void {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+    // send team_create_request message
+    try {
+      this.connection.send({
+        type: "control_team_create",
+        client_message_id: uuidv4(),
+        team_name: teamName,
+        initial_members: initialMembers,
+      });
+    } catch (error) {
+      this.eventHandler.onError(String(error));
+    }
+  }
+
+  /**
+   * Add a user to a team
+   * @param teamId - UUID of the team
+   * @param userId - UUID of the user to add
+   */
+  addTeamMember(teamId: string, userId: string): void {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+    try {
+      this.connection.send({
+        type: "control_add_team_user",
+        client_message_id: uuidv4(),
+        target_team_id: teamId,
+        user_uuid_or_email: userId,
+      });
+    } catch (error) {
+      this.eventHandler.onError(String(error));
+    }
+  }
+
+  /**
+   * Remove a user from a team
+   * @param teamId - UUID of the team
+   * @param userId - UUID of the user to remove
+   */
+  removeTeamMember(teamId: string, userId: string): void {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+    try {
+      this.connection.send({
+        type: "control_remove_team_user",
+        client_message_id: uuidv4(),
+        target_team_id: teamId,
+        user_uuid_or_email: userId,
+      });
+    } catch (error) {
+      this.eventHandler.onError(String(error));
+    }
+  }
+
+  /**
+   * Create a new agent
+   * @param agentName - the agent name
+   * @param templateName - the template name
+   * @param teamUuid - the team uuid
+   */
+  createNewAgent(
+    agentName: string,
+    templateName?: string,
+    teamUuid?: string
+  ): void {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+    try {
+      this.connection.send({
+        type: "control_agent_profile_create",
+        title: agentName,
+        template_name: templateName,
+        team_uuid: teamUuid,
+      });
+    } catch (error) {
+      this.eventHandler.onError(String(error));
+    }
+  }
+
+  /**
+   * Delete a session from the agent session map
+   * @param agentSessionMap - the agent session map
+   * @param agentUuid - the agent uuid
+   * @param sessionId - the session id
+   * @returns true if the session is deleted, false otherwise
+   */
+  private static deleteSessionFromAgentSessionMap(
+    agentSessionMap: Map<string, AgentSessionData>,
+    agentUuid: string,
+    sessionId: string
+  ): boolean {
+    const agentSession = agentSessionMap.get(agentUuid);
+    if (!agentSession) {
+      return false;
+    }
+    const size = agentSession.sessions.length;
+    agentSession.sessions = agentSession.sessions.filter(
+      (session) => session.session_uuid !== sessionId
+    );
+    return size > agentSession.sessions.length;
+  }
+
+  private handleSessionDeleted(msg: ServerControlSessionDeleted): void {
+    const sessionId = msg.session_id;
+    const agentUuid = msg.agent_profile_uuid;
+    logger.info(`[ChatClient] Session deleted: ${JSON.stringify(msg)}`);
+    let deleted = false;
+    if (msg.team_uuid) {
+      deleted = ChatClient.deleteSessionFromAgentSessionMap(
+        this.teams.get(msg.team_uuid)?.agentSessionMap ??
+          new Map<string, AgentSessionData>(),
+        agentUuid,
+        sessionId
+      );
+    } else {
+      deleted = ChatClient.deleteSessionFromAgentSessionMap(
+        this.userAgentSessionMap,
+        agentUuid,
+        sessionId
+      );
+    }
+    if (!deleted) {
+      // could be a race condition, we fetch the session list from the server
+      logger.warn(`[ChatClient] Session ${sessionId} is not in session map`);
+      this.connection.send({
+        type: "control_get_session_list",
+        client_message_id: uuidv4(),
+      });
+    }
+    this.activeSessions.delete(sessionId);
+    this.newMessage.delete(sessionId);
+    void this.eventHandler.onMessage(msg, this);
+  }
+
+  private handleControlMessage(msg: ServerControlMessage): void {
+    switch (msg.type) {
+      case "control_session_left":
+        this.activeSessions.delete(msg.session_id);
+        break;
+      case "control_agent_profile_created":
+        this.handleAgentProfileCreated(msg);
+        break;
+      case "control_agent_profile_deleted":
+        this.handleAgentProfileDeleted(msg);
+        break;
+      case "control_error":
+        this.handleControlError(msg);
+        break;
+      case "control_session_deleted":
+        this.handleSessionDeleted(msg);
+        break;
+      case "control_session_list":
+        logger.warn(`[ChatClient] Unexpected message in runtime: ${msg.type} `);
+        break;
+      case "control_team_created":
+        this.handleTeamCreatedMessage(msg);
+        break;
+      default: {
+        const _exhaustive: never = msg;
+        throw new Error(`unexpected control msg: ${JSON.stringify(msg)}`);
+      }
+    }
+    void this.eventHandler.onMessage(msg, this);
+  }
+
+  private handleAgentProfileCreated(msg: ServerControlAgentProfileCreated) {
+    logger.debug(
+      `[ChatClient.handleAgentProfileCreated] msg: ${JSON.stringify(msg)}`
+    );
+
+    const teamId = msg.profile.team_uuid;
+    const profileId = msg.profile.uuid;
+    const newAgentSession = {
+      agent_profile: msg.profile,
+      sessions: [],
+      updated_at: Date.now(),
+    };
+    if (teamId) {
+      const team = this.teams.get(teamId);
+      if (!team) {
+        throw new Error(`[handleAgentProfileCreated] no such team ${teamId}`);
+      }
+      team.agentSessionMap.set(profileId, newAgentSession);
+    } else {
+      this.userAgentSessionMap.set(profileId, newAgentSession);
+    }
+    void this.eventHandler.onMessage(msg, this);
+  }
+
+  private handleAgentProfileDeleted(_msg: ServerControlAgentProfileDeleted) {
+    throw new Error("handleAgentProfileDeleted not implemented");
+  }
+
+  private handleSessionMessage(msg: ServerSessionScopedMessage): void {
+    const sessionId = msg.session_id;
+
+    // TODO: we need to handle create new session case here, previously create
+    // new session's param are passed by URL, this is no longer an option
+    // since we are not creating a new connection for each session.
+    // Handle session_info message for pending session joins
+    if (msg.type === "session_info") {
+      this.handleSessionInfoMessage(msg);
+      // notify the event handler, i.e. UI that the session info update
+      void this.eventHandler.onMessage(msg, this);
+      return;
+    }
+
+    // Handle messages for existing active sessions
+    if (this.activeSessions.has(sessionId)) {
+      const sessionClient = this.activeSessions.get(sessionId);
+      if (sessionClient) {
+        sessionClient.handleMessage(msg);
+        // we are only calling the onMessage for session scoped message for now.
+        void this.eventHandler.onMessage(msg, this);
+      } else {
+        throw new Error(`Session client not found for session ${sessionId}`);
+      }
+    } else {
+      // Session not active yet, mark as having new messages
+      this.newMessage.set(sessionId, true);
+    }
+  }
+
+  private handleTeamCreatedMessage(msg: ServerControlTeamCreated): void {
+    const teamInfo: TeamInfo = {
+      team_uuid: msg.team_uuid,
+      team_name: msg.team_name,
+      owner_uuid: msg.team_owner_uuid,
+      participants: msg.members,
+      sessions: [],
+      agents: [],
+    };
+    this.teams.set(teamInfo.team_uuid, ChatClient.toClientTeamInfo(teamInfo));
+    this.currentTeamId = teamInfo.team_uuid;
+    void this.eventHandler.onMessage(msg, this);
+    return;
+  }
+
+  /**
+   * Handle session_info message which can be for:
+   * 1. A pending session join/create request
+   * 2. An update to an existing active session
+   * We update the session list here as well.
+   */
+  private handleSessionInfoMessage(msg: ServerSessionInfo): void {
+    const sessionId = msg.session_id;
+
+    // There is a pending request, check if this is the response
+    if (
+      this.sessionJoinOrCreateRes &&
+      msg.client_message_id === this.sessionJoinOrCreateRes.clientMessageId
+    ) {
+      const { resolve, reject } = this.sessionJoinOrCreateRes;
+
+      try {
+        logger.info(
+          `[ChatClient] Creating SessionClient for session ${sessionId}`
+        );
+        logger.info(`[ChatClient] msg: ${JSON.stringify(msg)}`);
+        const sessionClient = new SessionClient(
+          sessionId,
+          msg.saved_agent_profile,
+          this.connection,
+          msg.mcp_server_briefs,
+          createSessionParticipantMap(msg.participants)
+        );
+
+        // we need to pass the session id if this is a new session
+        if (this.sessionJoinOrCreateRes.sessionId === "") {
+          this.sessionJoinOrCreateRes.sessionId = sessionId;
+        } else if (this.sessionJoinOrCreateRes.sessionId !== sessionId) {
+          throw new Error(
+            `[ChatClient] session id mismatch: ` +
+              `${this.sessionJoinOrCreateRes.sessionId} !== ${sessionId}`
+          );
+        }
+
+        resolve(sessionClient, msg);
+      } catch (error) {
+        const errorMsg =
+          error instanceof Error ? error : new Error(String(error));
+        logger.error(
+          `[ChatClient] Failed to create SessionClient: ${errorMsg.message}`
+        );
+        reject(errorMsg);
+      }
+    } else {
+      this.updateSessionInfo(msg);
+    }
+  }
+
+  handleControlError(msg: ServerControlError): void {
+    // reject the pending session join/create request if message id matches
+    if (
+      this.sessionJoinOrCreateRes &&
+      msg.client_message_id === this.sessionJoinOrCreateRes.clientMessageId
+    ) {
+      this.sessionJoinOrCreateRes.reject(new Error(msg.message));
+    }
+    this.eventHandler.onError(`Server error: ${msg.message}`);
+  }
+
+  /**
+   * Update the session list. This is called when a session is updated.
+   * An error is thrown if the session id is not found.
+   * TODO: we might want to resync session map, for now, throw error to
+   * expose problems in the code.
+   */
+  private updateAgentSessionMap(sessionInfo: ServerSessionInfo): void {
+    const sessionId = sessionInfo.session_id;
+    if (sessionInfo.team_uuid) {
+      const teamInfo = this.teams.get(sessionInfo.team_uuid);
+      if (!teamInfo) {
+        throw new Error(`Team ${sessionInfo.team_uuid} not found in team list`);
+      }
+      //teamInfo.sessions.set(sessionId, ChatClient.toSessionData(sessionInfo));
+      const agentSessionMap = teamInfo.agentSessionMap;
+      if (!doUpdateAgentSessionMap(agentSessionMap, sessionInfo)) {
+        throw new Error(
+          `[updateAgentSessionMap] team session ${sessionId}` +
+            ` not found in agent session map`
+        );
+      }
+    } else {
+      const agentSessionMap = this.userAgentSessionMap;
+      if (!doUpdateAgentSessionMap(agentSessionMap, sessionInfo)) {
+        throw new Error(
+          `[updateAgentSessionMap] user session ${sessionId}` +
+            ` not found in agent session map`
+        );
+      }
+    }
+  }
+
+  /**
+   * Update the session info for an existing session. This passes the session
+   * info to the session client. An error is thrown if the session client is not
+   * found.
+   */
+  private updateSessionInfo(msg: ServerSessionInfo): void {
+    const sessionId = msg.session_id;
+    const sessionClient = this.activeSessions.get(sessionId);
+    this.updateAgentSessionMap(msg);
+    if (sessionClient) {
+      sessionClient.updateSessionInfo(msg);
+    } else {
+      throw new Error(`Session client not found for session ${sessionId}`);
+    }
+  }
+
+  private addSessionToAgentSessionMap(msg: ServerSessionInfo): void {
+    // get the correct agent session map
+    const agentSessionMap = msg.team_uuid
+      ? this.teams.get(msg.team_uuid)?.agentSessionMap
+      : this.userAgentSessionMap;
+    if (!agentSessionMap) {
+      throw new Error(`Invalid team uuid ${JSON.stringify(msg.team_uuid)}`);
+    }
+    // update the agent session map
+    const agentUuid = msg.saved_agent_profile.uuid;
+    const agentSession = agentSessionMap.get(agentUuid);
+    if (!agentSession) {
+      throw new Error(
+        `[addSessionToAgentSessionMap] Agent ${agentUuid}` +
+          ` not found in agent session map`
+      );
+    }
+    agentSession.sessions.push(sessionInfoToSessionData(msg));
+    agentSession.updated_at = new Date(msg.updated_at).getTime();
+  }
+}
+
+function sessionInfoToSessionData(msg: ServerSessionInfo): SessionData {
+  return {
+    session_uuid: msg.session_id,
+    title: msg.title,
+    team_uuid: msg.team_uuid,
+    agent_profile_uuid: msg.saved_agent_profile.uuid,
+    workspace: msg.workspace,
+    updated_at: msg.updated_at,
+    user_uuid: msg.owner_uuid,
+  };
+}
+
+/**
+ * Update the agent session map
+ * @param agentSessionMap - the agent session map
+ * @param sessionInfo - the session info
+ * @returns true if the session is updated, false otherwise
+ */
+function doUpdateAgentSessionMap(
+  agentSessionMap: Map<string, AgentSessionData>,
+  sessionInfo: ServerSessionInfo
+): boolean {
+  const sessionId = sessionInfo.session_id;
+  const agentUuid = sessionInfo.saved_agent_profile.uuid;
+  const agent = agentSessionMap.get(agentUuid);
+  if (!agent) {
+    return false;
+  }
+  let updated = false;
+  agent.sessions.map((session) => {
+    if (session.session_uuid === sessionId) {
+      updated = true;
+      return sessionInfoToSessionData(sessionInfo);
+    } else {
+      return session;
+    }
+  });
+  return updated;
+}