@xalia/agent 0.6.1 → 0.6.3

package/src/chat/client/chatClient.ts

@@ -1,4 +1,5 @@
 import { v4 as uuidv4 } from "uuid";
+import { strict as assert } from "assert";
 
 import { getLogger, SavedAgentProfile } from "@xalia/xmcp/sdk";
 
@@ -6,10 +7,10 @@ import { Connection } from "./connection";
 import { SessionClient } from "./sessionClient";
 import { IChatClientEventHandler } from "./interfaces";
 import {
-  SessionData,
   TeamInfo,
   TeamParticipant,
   AgentSessionData,
+  SessionDescriptor,
 } from "../data/dataModels";
 import { createSessionParticipantMap } from "../data/database";
 import {
@@ -21,17 +22,34 @@ import {
   ServerControlMessage,
   ServerControlSessionDeleted,
   ServerControlTeamCreated,
+  ServerControlTeamMembersUpdated,
   ServerToClient,
   ClientToServer,
   ServerControlAgentProfileCreated,
   ServerControlAgentProfileDeleted,
+  ClientControlSessionCreate,
 } from "../protocol/messages";
-import { SESSION_CREATE_TIMEOUT, SESSION_JOIN_TIMEOUT } from "./constants";
+import { SESSION_JOIN_TIMEOUT } from "./constants";
 import { ITeamManager } from "./teamManager";
 import { buildAgentSessionMap } from "../utils/agentSessionMap";
+import { ResponseAwaiter } from "../utils/responseAwaiter";
 
 const logger = getLogger();
 
+/// A fake message which includes an actual client.  This avoids a peculiarity
+/// related to handling of async resolves.
+
+type CreatedClient = {
+  client: SessionClient;
+  sessionInfo: ServerSessionInfo;
+};
+
+type JoinCreateResponse = ServerSessionInfo | ServerControlError;
+
+function idExtractor(msg: JoinCreateResponse): string | undefined {
+  return msg.client_message_id;
+}
+
 // client team info, used for session tracking etc.
 export type ClientTeamInfo = {
   team_uuid: string;
@@ -41,37 +59,54 @@ export type ClientTeamInfo = {
   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 connection: Connection<ClientToServer, ServerToClient>;
+  private eventHandler: IChatClientEventHandler;
+  // agent_uuid -> agent session data
+  private userAgentSessionMap: Map<string, AgentSessionData>;
+  // team_uuid -> team info
+  private teams: Map<string, ClientTeamInfo>;
+  // Whether this client is in guest mode
+  private isGuestMode: boolean;
+  private closed: boolean;
+  private currentSessionId: string | undefined;
+  // note: currentTeamId will not be reset when swtiching to a user session.
+  private currentTeamId: string | undefined;
+  // session_uuid -> session client
+  private activeSessions: Map<string, SessionClient>;
+  // session_uuid -> true if there is a new message.
+  private newMessage: Map<string, boolean>;
+  // Track pending session join requests.
+  private sessionJoinAwaiter: ResponseAwaiter<
+    JoinCreateResponse,
+    CreatedClient
+  >;
+
   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
-  ) {}
+    connection: Connection<ClientToServer, ServerToClient>,
+    eventHandler: IChatClientEventHandler,
+    userAgentSessionMap: Map<string, AgentSessionData>,
+    teams: Map<string, ClientTeamInfo>,
+    isGuestMode: boolean
+  ) {
+    this.connection = connection;
+    this.eventHandler = eventHandler;
+    this.userAgentSessionMap = userAgentSessionMap;
+    this.teams = teams;
+    this.isGuestMode = isGuestMode;
+
+    this.closed = false;
+    this.currentSessionId = undefined;
+    this.currentTeamId = undefined;
+    this.activeSessions = new Map();
+    this.newMessage = new Map();
+    this.sessionJoinAwaiter = ResponseAwaiter.initWithImmediate(
+      "control_error",
+      idExtractor,
+      this.onJoinCreateResponse.bind(this),
+      Math.max(SESSION_JOIN_TIMEOUT)
+    );
+  }
 
   /**
    * Connect to server and create ChatClient instance
@@ -81,6 +116,9 @@ export class ChatClient implements ITeamManager {
     token: string,
     eventHandler: IChatClientEventHandler
   ): Promise<ChatClient> {
+    // Determine if this is a guest token (format: guest_<hash>_<session-id>)
+    const isGuestMode = token.startsWith("guest_");
+
     const connection = new Connection<ClientToServer, ServerToClient>({
       url,
       token,
@@ -92,7 +130,7 @@ export class ChatClient implements ITeamManager {
       // 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>();
+        const userSessions = new Map<string, SessionDescriptor>();
         msg.user_sessions.forEach((session) => {
           userSessions.set(session.session_uuid, session);
         });
@@ -127,7 +165,8 @@ export class ChatClient implements ITeamManager {
             connection,
             eventHandler,
             userAgentSessionMap,
-            teams
+            teams,
+            isGuestMode
           );
           resolveClient(client);
         } else {
@@ -291,7 +330,6 @@ export class ChatClient implements ITeamManager {
     this.teams.clear();
     this.currentSessionId = undefined;
     this.currentTeamId = undefined;
-    this.sessionJoinOrCreateRes = undefined;
   }
 
   public async createNewSession(
@@ -303,74 +341,24 @@ export class ChatClient implements ITeamManager {
       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");
-    }
+    const clientMessageId = uuidv4();
+    const createdSessionP =
+      this.sessionJoinAwaiter.waitForResponse(clientMessageId);
 
-    return new Promise((resolve, reject) => {
-      const clientMessageId = uuidv4();
+    const createMsg: ClientControlSessionCreate = {
+      type: "control_session_create",
+      client_message_id: clientMessageId,
+      title: newTitle,
+      agent_profile_id: agentProfileId,
+      team_id: teamId,
+    };
+    this.connection.send(createMsg);
+    logger.info(`[ChatClient] Sent ${JSON.stringify(createMsg)}`);
 
-      // 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)));
-      }
-    });
+    const createdSession = await createdSessionP;
+    const sessionId = createdSession.client.getSessionUUID();
+    logger.info(`[ChatClient] joined session ${sessionId}`);
+    return createdSession.client;
   }
 
   /**
@@ -382,66 +370,26 @@ export class ChatClient implements ITeamManager {
       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();
+    const clientMessageId = uuidv4();
+    const createdSessionP =
+      this.sessionJoinAwaiter.waitForResponse(clientMessageId);
 
-      // 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)));
-      }
+    this.connection.send({
+      type: "control_session_join",
+      client_message_id: clientMessageId,
+      target_session_id: sessionId,
     });
+    logger.debug(`[ChatClient] Sent session_join for session ${sessionId}`);
+
+    // Await the response
+
+    const createdSession = await createdSessionP;
+    const newSessionId = createdSession.client.getSessionUUID();
+    if (newSessionId !== sessionId) {
+      throw new Error(`unexpected session id ${newSessionId}`);
+    }
+    logger.info(`[ChatClient] joined session ${sessionId}`);
+    return createdSession.client;
   }
 
   /**
@@ -476,6 +424,27 @@ export class ChatClient implements ITeamManager {
     );
   }
 
+  /**
+   * Delete an agent profile by sending control_agent_profile_delete message
+   * @param agentProfileUuid - The UUID of the agent profile to delete
+   * @returns void
+   */
+  deleteAgentProfile(agentProfileUuid: string): void {
+    if (this.closed) {
+      throw new Error("ChatClient is closed");
+    }
+
+    this.connection.send({
+      type: "control_agent_profile_delete",
+      agent_profile_uuid: agentProfileUuid,
+    });
+
+    logger.debug(
+      `[ChatClient] Sent control_agent_profile_delete for profile ` +
+        agentProfileUuid
+    );
+  }
+
   /**
    * Create a new team with initial members
    * @param teamName
@@ -644,6 +613,9 @@ export class ChatClient implements ITeamManager {
       case "control_team_created":
         this.handleTeamCreatedMessage(msg);
         break;
+      case "control_team_members_updated":
+        this.handleTeamMembersUpdated(msg);
+        break;
       default: {
         const _exhaustive: never = msg;
         throw new Error(`unexpected control msg: ${JSON.stringify(msg)}`);
@@ -676,8 +648,36 @@ export class ChatClient implements ITeamManager {
     void this.eventHandler.onMessage(msg, this);
   }
 
-  private handleAgentProfileDeleted(_msg: ServerControlAgentProfileDeleted) {
-    throw new Error("handleAgentProfileDeleted not implemented");
+  private handleAgentProfileDeleted(msg: ServerControlAgentProfileDeleted) {
+    logger.debug(
+      `[ChatClient.handleAgentProfileDeleted] msg: ${JSON.stringify(msg)}`
+    );
+
+    const profileUuid = msg.profile_uuid;
+    let found = false;
+
+    // Try to remove from user agent session map
+    if (this.userAgentSessionMap.has(profileUuid)) {
+      this.userAgentSessionMap.delete(profileUuid);
+      found = true;
+    } else {
+      // Try to remove from team agent session maps
+      for (const [_teamId, team] of this.teams.entries()) {
+        if (team.agentSessionMap.has(profileUuid)) {
+          team.agentSessionMap.delete(profileUuid);
+          found = true;
+          break;
+        }
+      }
+    }
+
+    if (!found) {
+      logger.warn(
+        `[ChatClient] Agent profile ${profileUuid} not found in any session map`
+      );
+    }
+
+    void this.eventHandler.onMessage(msg, this);
   }
 
   private handleSessionMessage(msg: ServerSessionScopedMessage): void {
@@ -725,119 +725,71 @@ export class ChatClient implements ITeamManager {
     return;
   }
 
+  private handleTeamMembersUpdated(msg: ServerControlTeamMembersUpdated): void {
+    const team = this.teams.get(msg.team_uuid);
+    if (!team) {
+      logger.warn(
+        `[ChatClient] Received team_members_updated for unknown team ` +
+          msg.team_uuid
+      );
+      return;
+    }
+
+    // Update the participants map
+    team.participants = new Map(
+      msg.members.map((participant) => [participant.user_uuid, participant])
+    );
+
+    logger.info(
+      `[ChatClient] Updated team members for team ${msg.team_uuid}, ` +
+        `now has ${String(msg.members.length)} members`
+    );
+    void this.eventHandler.onMessage(msg, this);
+  }
+
   /**
    * Handle session_info message which can be for:
    * 1. A pending session join/create request
    * 2. An update to an existing active session
+   * 3. A new session notification
    * 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}`
-          );
-        }
+    // We intercept messages for th Awaiter here so that we can create and
+    // register the client immediately, before other messages for this session
+    // are handled.  Relying on Awaiter alone (currently) schedules the resolv
+    // to be called later (potentially AFTER further messages for the session
+    // have been received).
 
-        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);
+    if (this.sessionJoinAwaiter.onMessage(msg)) {
+      return;
     }
+
+    logger.debug(
+      `[ChatClient.handleSessionInfoMessage] not handled by awaiter - updating`
+    );
+    this.upsertSessionToAgentSessionMap(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));
+    if (this.sessionJoinAwaiter.onMessage(msg)) {
+      return;
     }
-    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`
-        );
-      }
-    }
+    this.eventHandler.onError(`Server error: ${msg.message}`);
   }
 
-  /**
-   * 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 upsertSessionToAgentSessionMap(msg: ServerSessionInfo): void {
+    // Skip agent session map updates for guest sessions since they don't
+    // need session organization
+    if (this.isGuestMode) {
+      logger.info(
+        `[ChatClient] Skipping agent session map add for ` +
+          `guest session ${msg.session_id}`
+      );
+      return;
     }
-  }
 
-  private addSessionToAgentSessionMap(msg: ServerSessionInfo): void {
     // get the correct agent session map
     const agentSessionMap = msg.team_uuid
       ? this.teams.get(msg.team_uuid)?.agentSessionMap
@@ -854,47 +806,71 @@ export class ChatClient implements ITeamManager {
           ` not found in agent session map`
       );
     }
-    agentSession.sessions.push(sessionInfoToSessionData(msg));
+
+    // Check if session already exists
+    const existingSessionIndex = agentSession.sessions.findIndex(
+      (session) => session.session_uuid === msg.session_id
+    );
+
+    if (existingSessionIndex !== -1) {
+      // Update existing session
+      agentSession.sessions[existingSessionIndex] =
+        sessionInfoToSessionDescriptor(msg);
+      logger.info(
+        `[ChatClient] Updated existing session ${msg.session_id} in ` +
+          `agent session map`
+      );
+    } else {
+      // Add new session
+      agentSession.sessions.push(sessionInfoToSessionDescriptor(msg));
+      logger.info(
+        `[ChatClient] Added new session ${msg.session_id} to agent session map`
+      );
+    }
+
     agentSession.updated_at = new Date(msg.updated_at).getTime();
   }
+
+  // This is the immediate callback of the ResponseAwaiter for session
+  // joining.  It must create and register the client.
+  private onJoinCreateResponse(msg: JoinCreateResponse): CreatedClient {
+    // Errors should not get through here.
+    assert(msg.type === "session_info");
+
+    const sessionId = msg.session_id;
+    const client = new SessionClient(
+      sessionId,
+      msg.saved_agent_profile,
+      this.connection,
+      msg.mcp_server_briefs,
+      createSessionParticipantMap(msg.participants),
+      msg.agent_paused
+    );
+
+    this.activeSessions.set(sessionId, client);
+    this.currentSessionId = sessionId;
+    if (msg.team_uuid) {
+      this.currentTeamId = msg.team_uuid;
+    }
+    this.upsertSessionToAgentSessionMap(msg);
+
+    return {
+      client,
+      sessionInfo: msg,
+    };
+  }
 }
 
-function sessionInfoToSessionData(msg: ServerSessionInfo): SessionData {
+function sessionInfoToSessionDescriptor(
+  msg: ServerSessionInfo
+): SessionDescriptor {
   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,
+    agent_paused: msg.agent_paused,
   };
 }
-
-/**
- * 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;
-}

package/src/chat/client/constants.ts

@@ -1,3 +1 @@
 export const SESSION_JOIN_TIMEOUT = 6000; // 6 seconds
-
-export const SESSION_CREATE_TIMEOUT = 6000; // 6 seconds