@xalia/agent 0.6.0 → 0.6.2

package/src/chat/server/sessionRegistry.ts

@@ -6,7 +6,6 @@ import {
   SavedAgentProfile,
   DEFAULT_AGENT_PROFILE_NAME,
   DEFAULT_AGENT_PROFILE,
-  DEFAULT_AGENT_PROFILE_SYSTEM_PROMPT,
 } from "@xalia/xmcp/sdk";
 
 import {
@@ -27,24 +26,28 @@ import {
   ClientControlRemoveTeamUser,
   ServerToClient,
   ClientControlAgentProfileCreate,
+  ClientControlAgentProfileDelete,
 } from "../protocol/messages";
-import { OpenSession } from "./openSession";
+import { GUEST_TOKEN_PREFIX, OpenSession } from "./openSession";
 import {
   SessionCreateData,
   SessionData,
+  SessionDescriptor,
   TeamInfo,
   TeamParticipant,
 } from "../data/dataModels";
-import { Database } from "../data/database";
+import { Database, UserData } from "../data/database";
 import { resolveUserIdentifier } from "../utils/userResolver";
 import { ChatFatalError } from "../protocol/errors";
-import { IUserConnectionManager } from "./connectionManager";
-import { DEFAULT_LLM_MODEL } from "../../agent/agentUtils";
+import { IMessageProcessor, IUserConnectionManager } from "./connectionManager";
 import { getErrorString } from "./errorUtils";
+import { ApiKeyManager } from "../data/apiKeyManager";
 
 const logger = getLogger();
 
-export class SessionRegistry {
+export type GuestUser = UserData & { guest_for_session: string };
+
+export class SessionRegistry implements IMessageProcessor<ClientToServer> {
   // In memory session-user/user-session mappings
   // Note: this mappings ONLY trackes online users and will
   // be cleaned up when the user disconnects.
@@ -57,12 +60,18 @@ export class SessionRegistry {
   // sessionId -> OpenSession
   private openSessions: Map<string, OpenSession> = new Map();
 
+  private guests: Map<string, GuestUser> = new Map();
+
+  private apiKeyManager: ApiKeyManager;
+
   constructor(
     private db: Database,
     private connectionManager: IUserConnectionManager<ServerToClient>,
     private llmUrl: string,
     private xmcpUrl: string
-  ) {}
+  ) {
+    this.apiKeyManager = new ApiKeyManager(db);
+  }
 
   /**
    * Add user to session membership.
@@ -168,6 +177,60 @@ export class SessionRegistry {
     return this.userSessions.get(userId) || new Set();
   }
 
+  // IMessageProcessor<ClientToServer>
+  async authenticate(token: string): Promise<string | undefined> {
+    // Parse the api key to determine the type of connection
+    const {
+      prefix,
+      apiKey,
+      payload: sessionUUID,
+    } = ApiKeyManager.parseToken(token);
+
+    if (prefix === GUEST_TOKEN_PREFIX) {
+      if (!sessionUUID) {
+        throw new Error(`invalid token ${token}`);
+      }
+
+      const descriptor = await this.getSessionDescriptor(sessionUUID);
+      if (!descriptor) {
+        throw new Error(`no such session ${sessionUUID}`);
+      }
+      if (descriptor.access_token !== token) {
+        throw new Error(`invalid guest key ${apiKey}`);
+      }
+
+      // For now, add a new user to the DB (so simplify the process of adding
+      // users).
+
+      const user_uuid = uuidv4();
+      const email = `${user_uuid}@`;
+      const nickname = "Guest";
+      const timezone = "UTC";
+
+      await this.db.createUser(user_uuid, email, nickname, timezone);
+
+      const user: GuestUser = {
+        uuid: user_uuid,
+        nickname: "Guest",
+        email: "guest",
+        timezone: "UTC",
+        guest_for_session: sessionUUID,
+      };
+      this.guests.set(user_uuid, user);
+      return user.uuid;
+
+      // throw new Error("unimpl");
+      // // return "";
+    }
+
+    const userData = await this.apiKeyManager.verifyApiKey(apiKey);
+    if (!userData) {
+      throw new ChatFatalError("invalid api key");
+    }
+    return userData.uuid;
+  }
+
+  // IMessageProcessor<ClientToServer>
   async processMessage(
     connectionId: string,
     userId: string,
@@ -204,11 +267,11 @@ export class SessionRegistry {
   ): Promise<void> {
     switch (message.type) {
       case "control_agent_profile_create":
-        await this.handleAgentProfileCreate(message, userId);
+        await this.handleAgentProfileCreate(userId, message);
         break;
       case "control_agent_profile_delete":
-        // TODO:
-        throw new Error("not implemented yet");
+        await this.handleAgentProfileDelete(userId, message);
+        break;
       case "control_get_session_list":
         await this.handleGetSessionList(connectionId, userId, message);
         break;
@@ -251,6 +314,19 @@ export class SessionRegistry {
     try {
       const [userSessions, teamSessions, userAgents] =
         await this.getAllAgentsAndSessionsByUser(userId);
+
+      // Mark any guest session as a user session
+
+      const guest = this.guests.get(userId);
+      if (guest) {
+        const guestSession = await this.db.sessionGetDescriptorById(
+          guest.guest_for_session
+        );
+        if (guestSession) {
+          userSessions.push(guestSession);
+        }
+      }
+
       const response: ServerControlSessionList = {
         type: "control_session_list",
         user_sessions: userSessions,
@@ -286,11 +362,11 @@ export class SessionRegistry {
       await this.validateSessionAccess(sessionId, userId, "owner");
 
       // get users/team-uuid before deletion
-      const [users, teamUuid, sessionData] = await Promise.all([
+      const [users, sessionData] = await Promise.all([
         this.getSessionUsers(sessionId),
-        this.db.sessionGetTeamUuid(sessionId),
-        this.db.sessionGetById(sessionId),
+        this.db.sessionGetDescriptorById(sessionId),
       ]);
+      const teamUuid = sessionData?.team_uuid;
 
       if (!sessionData) {
         throw new ChatFatalError(`No such session: ${sessionId}`);
@@ -355,9 +431,10 @@ export class SessionRegistry {
       const validMemberIds: string[] = [];
       const participants: TeamParticipant[] = [];
 
-      // filter out the owner and extract failed lookups
+      // No need to filter out the owner here, that is done in
+      // `createTeamWithParticipants`
       resolvedMemberIds.forEach((user, index) => {
-        if (user && user.uuid !== userId) {
+        if (user) {
           validMemberIds.push(user.uuid);
           participants.push({
             user_uuid: user.uuid,
@@ -365,7 +442,7 @@ export class SessionRegistry {
             email: user.email,
             role: "participant",
           });
-        } else if (user === undefined) {
+        } else {
           failedLookups.push(message.initial_members[index]);
         }
       });
@@ -377,17 +454,50 @@ export class SessionRegistry {
         validMemberIds
       );
 
+      // Create a default agent for the new team
+      const defaultAgentProfile = await this.db.createAgentProfile(
+        undefined,
+        teamUuid,
+        DEFAULT_AGENT_PROFILE_NAME,
+        DEFAULT_AGENT_PROFILE
+      );
+
+      if (!defaultAgentProfile) {
+        throw new Error(
+          `Failed to create default agent profile for team ${teamUuid}`
+        );
+      }
+
+      // Get owner's information and add to participants for the response
+      const ownerInfo = await this.db.getUserFromUuid(userId);
+      if (!ownerInfo) {
+        throw new Error(`Cannot find owner user data for user ${userId}`);
+      }
+
+      const ownerParticipant: TeamParticipant = {
+        user_uuid: userId,
+        nickname: ownerInfo.nickname || "",
+        email: ownerInfo.email,
+        role: "owner",
+      };
+
       const response: ServerControlTeamCreated = {
         type: "control_team_created",
         team_uuid: teamUuid,
         team_owner_uuid: userId,
         team_name: message.team_name,
-        members: participants,
+        members: [ownerParticipant, ...participants],
         failed_lookups: failedLookups,
       };
       const members = new Set(validMemberIds);
       members.add(userId);
       this.connectionManager.sendToUsers(members, response);
+
+      // Broadcast the default agent profile to all team members
+      this.connectionManager.sendToUsers(members, {
+        type: "control_agent_profile_created",
+        profile: defaultAgentProfile,
+      });
     } catch (error) {
       logger.error(`[SessionRegistry] Error creating team: ${String(error)}`);
       this.sendControlError(
@@ -500,6 +610,19 @@ export class SessionRegistry {
       }
     }
 
+    // Notify all team members about the updated member list
+    const updatedParticipants = await this.db.teamGetMembers(
+      message.target_team_id
+    );
+    this.connectionManager.sendToUsers(
+      new Set(updatedParticipants.map((p) => p.user_uuid)),
+      {
+        type: "control_team_members_updated",
+        team_uuid: message.target_team_id,
+        members: updatedParticipants,
+      }
+    );
+
     logger.info(
       `[SessionRegistry] User ${user.uuid}` +
         `added to team ${message.target_team_id}`
@@ -592,6 +715,19 @@ export class SessionRegistry {
       }
     }
 
+    // Notify all remaining team members about the updated member list
+    const updatedParticipants = await this.db.teamGetMembers(
+      message.target_team_id
+    );
+    this.connectionManager.sendToUsers(
+      new Set(updatedParticipants.map((p) => p.user_uuid)),
+      {
+        type: "control_team_members_updated",
+        team_uuid: message.target_team_id,
+        members: updatedParticipants,
+      }
+    );
+
     logger.info(
       `[SessionRegistry] User ${user.uuid} ` +
         `removed from team ${message.target_team_id}`
@@ -645,7 +781,7 @@ export class SessionRegistry {
       const access = await this.validateSessionAccess(sessionId, userId);
       if (!access) {
         throw new ChatFatalError(
-          `User ${userId} is not authorized to ` + `join session ${sessionId}`
+          `User ${userId} is not authorized to join session ${sessionId}`
         );
       }
 
@@ -658,7 +794,15 @@ export class SessionRegistry {
           `Server internal error: ` + `failed to load session ${sessionId}`
         );
       }
-
+      const guest = this.guests.get(userId);
+      if (guest) {
+        session.addParticipant(userId, {
+          user_uuid: guest.uuid,
+          nickname: guest.nickname || "Guest",
+          email: guest.email,
+          role: "participant",
+        });
+      }
       // To this point, there should be no error thrown.
       // Update in-memory session-user/user-session mappings
       this.addUserToSessionMemory(userId, sessionId);
@@ -667,6 +811,7 @@ export class SessionRegistry {
       if (!this.openSessions.has(sessionId)) {
         this.openSessions.set(sessionId, session);
       }
+
       // pass the message to the session to handle the rest
       session.sendSessionData(connectionId, message.client_message_id);
     } catch (error) {
@@ -682,8 +827,8 @@ export class SessionRegistry {
   }
 
   async handleAgentProfileCreate(
-    message: ClientControlAgentProfileCreate,
-    from: string
+    userId: string,
+    message: ClientControlAgentProfileCreate
   ): Promise<string> {
     // get agent profile from template
     let agentProfileFromTemplate: AgentProfile | undefined = undefined;
@@ -697,11 +842,13 @@ export class SessionRegistry {
       agentProfileFromTemplate = template.agent_profile;
     }
 
-    const newAgentProfile: AgentProfile = agentProfileFromTemplate || {
-      model: DEFAULT_LLM_MODEL,
-      system_prompt: DEFAULT_AGENT_PROFILE_SYSTEM_PROMPT,
-      mcp_settings: {},
-    };
+    const newAgentProfile: AgentProfile =
+      agentProfileFromTemplate ||
+      new AgentProfile(
+        DEFAULT_AGENT_PROFILE.model,
+        DEFAULT_AGENT_PROFILE.system_prompt,
+        DEFAULT_AGENT_PROFILE.mcp_settings
+      );
     const team_uuid = message.team_uuid || undefined;
 
     if (team_uuid) {
@@ -732,7 +879,7 @@ export class SessionRegistry {
     // User AgentProfile
 
     const savedAgentProfile = await this.db.createAgentProfile(
-      from,
+      userId,
       undefined,
       message.title,
       newAgentProfile
@@ -742,7 +889,7 @@ export class SessionRegistry {
     }
 
     // Send the new AgentProfile to the user
-    this.connectionManager.sendToUsers(new Set([from]), {
+    this.connectionManager.sendToUsers(new Set([userId]), {
       type: "control_agent_profile_created",
       profile: savedAgentProfile,
     });
@@ -750,6 +897,54 @@ export class SessionRegistry {
     return savedAgentProfile.uuid;
   }
 
+  async handleAgentProfileDelete(
+    userId: string,
+    message: ClientControlAgentProfileDelete
+  ): Promise<void> {
+    const agentProfileUuid = message.agent_profile_uuid;
+
+    // Get the agent profile to determine if it's a team or user profile
+    const agentProfile =
+      await this.db.getSavedAgentProfileById(agentProfileUuid);
+    if (!agentProfile) {
+      throw new Error(`Agent profile ${agentProfileUuid} not found`);
+    }
+
+    // Validate team access: user must own the profile or be in the team
+    // TODO: Only allow the owner to delete the profile?
+    if (agentProfile.team_uuid) {
+      const hasAccess = await this.validateTeamAccess(
+        agentProfile.team_uuid,
+        userId
+      );
+      if (!hasAccess) {
+        throw new Error(
+          `User ${userId} is not authorized to delete this agent profile`
+        );
+      }
+    }
+
+    // Delete the agent profile from database
+    await this.db.deleteAgentProfile(agentProfileUuid);
+
+    // Send confirmation to all relevant users
+    if (agentProfile.team_uuid) {
+      const participants = await this.db.teamGetMembers(agentProfile.team_uuid);
+      this.connectionManager.sendToUsers(
+        new Set(participants.map((p) => p.user_uuid)),
+        {
+          type: "control_agent_profile_deleted",
+          profile_uuid: agentProfileUuid,
+        }
+      );
+    } else {
+      this.connectionManager.sendToUsers(new Set([userId]), {
+        type: "control_agent_profile_deleted",
+        profile_uuid: agentProfileUuid,
+      });
+    }
+  }
+
   /**
    * Create a new session for a user.
    * - create session in database (via `newSession`)
@@ -769,13 +964,13 @@ export class SessionRegistry {
 
         // Create AgentProfile and inform the client
         message.agent_profile_id = await this.handleAgentProfileCreate(
+          fromUserId,
           {
             type: "control_agent_profile_create",
             title: "New Agent " + uuidv4(),
             user_uuid: fromUserId,
             team_uuid: message.team_id,
-          },
-          fromUserId
+          }
         );
       }
 
@@ -825,8 +1020,8 @@ export class SessionRegistry {
     // validate the agent profile
     await this.validateSavedAgentProfile(agentProfileId);
 
-    const newSessionData: SessionData = {
-      ...userSessionDataCreate(ownerId, title, agentProfileId),
+    const newSessionData: SessionDescriptor = {
+      ...userSessionCreateData(ownerId, title, agentProfileId),
       updated_at: new Date().toISOString(),
     };
     const openSession = await OpenSession.initWithEmptySession(
@@ -860,8 +1055,8 @@ export class SessionRegistry {
       );
     }
 
-    const newSessionData: SessionData = {
-      ...teamSessionDataCreate(teamId, fromUserId, title, agentProfileId),
+    const newSessionData: SessionDescriptor = {
+      ...teamSessionCreateData(teamId, fromUserId, title, agentProfileId),
       updated_at: new Date().toISOString(),
     };
 
@@ -917,9 +1112,27 @@ export class SessionRegistry {
    * Handle user disconnect - clean up from all sessions.
    * Called when a connection is closed to ensure proper cleanup.
    */
-  handleUserDisconnect(userId: string): void {
+  async handleUserDisconnect(userId: string): Promise<void> {
     logger.info(`[SessionRegistry] Handling disconnect for user ${userId}`);
 
+    // If the user is a guest, remove his DB entry.
+    const guest = this.guests.get(userId);
+    if (guest) {
+      const session = this.openSessions.get(guest.guest_for_session);
+      if (session) {
+        session.removeParticipant(userId);
+      }
+      this.guests.delete(userId);
+      try {
+        logger.info(
+          `[SessionRegistry.handleUserDisconnect] deleting user ${userId}`
+        );
+        await this.db.deleteUser(userId);
+      } catch (e) {
+        logger.warn(`error removing guest ${userId}: ${getErrorString(e)}`);
+      }
+    }
+
     // Get all sessions the user is in (copy to avoid concurrent modification)
     const userSessionIds = this.getInMemoryUserSessions(userId);
     const sessionsToLeave = Array.from(userSessionIds);
@@ -943,11 +1156,13 @@ export class SessionRegistry {
   private async getAllAgentsAndSessionsByUser(
     userId: string
   ): Promise<[SessionData[], TeamInfo[], SavedAgentProfile[]]> {
-    return Promise.all([
+    const [userSessions, teamInfos, userAgents] = await Promise.all([
       this.db.getUserSessions(userId),
       this.db.getTeamInfosByUser(userId),
       this.agentProfilesGetByUserOrDefault(userId),
     ]);
+
+    return [userSessions, teamInfos, userAgents];
   }
 
   private async agentProfilesGetByUserOrDefault(
@@ -987,7 +1202,7 @@ export class SessionRegistry {
   /**
    * Validates that a user has permission to access a session.
    */
-  private async validateSessionAccess(
+  async validateSessionAccess(
     sessionId: string,
     userId: string,
     accessType?: "participant" | "owner"
@@ -998,6 +1213,13 @@ export class SessionRegistry {
       const participants = session.getParticipants();
       const role = participants.get(userId);
       if (!role) {
+        // Check for guest access
+        const guest = this.guests.get(userId);
+        if (guest && guest.guest_for_session === sessionId) {
+          return !accessType || accessType === "participant";
+        }
+
+        // Not authorized to join the session
         return false;
       } else {
         return !accessType || role.role === accessType;
@@ -1007,6 +1229,13 @@ export class SessionRegistry {
       const participants = await this.db.sessionGetParticipants(sessionId);
       const role = participants.find((p) => p.user_uuid === userId)?.role;
       if (!role) {
+        // Check for guest access
+        const guest = this.guests.get(userId);
+        logger.info(`[validateSessionAccess] guest: ${JSON.stringify(guest)}`);
+        if (guest && guest.guest_for_session === sessionId) {
+          return !accessType || accessType === "participant";
+        }
+
         return false;
       } else {
         return !accessType || role === accessType;
@@ -1030,9 +1259,24 @@ export class SessionRegistry {
       return !accessType || role === accessType;
     }
   }
+
+  /**
+   * Read from DB if the session is not active
+   */
+  private async getSessionDescriptor(
+    sessionId: string
+  ): Promise<SessionDescriptor | undefined> {
+    const session = this.openSessions.get(sessionId);
+    if (session) {
+      return session.getDescriptor();
+    } else {
+      // fetch the session from database
+      return this.db.sessionGetDescriptorById(sessionId);
+    }
+  }
 }
 
-export function userSessionDataCreate(
+export function userSessionCreateData(
   ownerId: string,
   title: string,
   agentProfileId: string
@@ -1042,12 +1286,12 @@ export function userSessionDataCreate(
     title,
     team_uuid: undefined,
     agent_profile_uuid: agentProfileId,
-    workspace: undefined,
     user_uuid: ownerId,
+    agent_paused: false,
   };
 }
 
-export function teamSessionDataCreate(
+export function teamSessionCreateData(
   teamId: string,
   ownerId: string,
   title: string,
@@ -1058,7 +1302,7 @@ export function teamSessionDataCreate(
     title,
     team_uuid: teamId,
     agent_profile_uuid: agentProfileId,
-    workspace: undefined,
     user_uuid: ownerId,
+    agent_paused: false,
   };
 }

package/src/chat/server/test-utils/mockFactories.ts

@@ -128,21 +128,18 @@ export function createMockUserConnectionManager(): {
   mock: IUserConnectionManager<ServerToClient>;
   spies: {
     sendToUsers: ReturnType<typeof vi.fn>;
-    getLiveUserApiKey: ReturnType<typeof vi.fn>;
   };
 } {
   const spies = {
     sendToUsers: vi.fn(),
     sendToConnection: vi.fn(),
     sendServerError: vi.fn(),
-    getLiveUserApiKey: vi.fn(),
   };
 
   const mock = {
     sendToUsers: spies.sendToUsers,
     sendToConnection: spies.sendToConnection,
     sendServerError: spies.sendServerError,
-    getLiveUserApiKey: spies.getLiveUserApiKey,
   } as IUserConnectionManager<ServerToClient>;
 
   return { mock, spies };
@@ -154,16 +151,25 @@ export function createMockUserConnectionManager(): {
 export function createMockSessionRegistry(): {
   mock: IMessageProcessor<ClientToServer>;
   spies: {
+    authenticate: ReturnType<typeof vi.fn>;
     processMessage: ReturnType<typeof vi.fn>;
     handleUserDisconnect: ReturnType<typeof vi.fn>;
   };
 } {
   const spies = {
+    authenticate: vi.fn().mockImplementation((apiKey) => {
+      if (apiKey === "valid-api-key")
+        return Promise.resolve(MOCK_USERS.owner.uuid);
+      if (apiKey === "participant-api-key")
+        return Promise.resolve(MOCK_USERS.participant.uuid);
+      return Promise.resolve(null);
+    }),
     processMessage: vi.fn(),
     handleUserDisconnect: vi.fn(),
   };
 
   const mock = {
+    authenticate: spies.authenticate,
     processMessage: spies.processMessage,
     handleUserDisconnect: spies.handleUserDisconnect,
   } as IMessageProcessor<ClientToServer>;
@@ -267,6 +273,7 @@ export function createMockSessionList(): Array<SessionData> {
       workspace: undefined,
       updated_at: MOCK_SESSIONS.active.updated_at || "",
       user_uuid: MOCK_SESSIONS.active.user_uuid,
+      agent_paused: false,
     },
     {
       session_uuid: MOCK_SESSIONS.secondary.uuid,
@@ -276,6 +283,7 @@ export function createMockSessionList(): Array<SessionData> {
       workspace: undefined,
       updated_at: MOCK_SESSIONS.secondary.updated_at || "",
       user_uuid: MOCK_SESSIONS.secondary.user_uuid,
+      agent_paused: false,
     },
   ];
 }
@@ -409,14 +417,6 @@ export function setupStandardMockBehaviors(mocks: {
   }
 
   // Setup user connection manager mocks
-  if (mocks.userConnectionManager) {
-    mocks.userConnectionManager.spies.getLiveUserApiKey.mockImplementation(
-      (userId: string) => {
-        if (userId === MOCK_USERS.owner.uuid) return "valid-api-key";
-        if (userId === MOCK_USERS.participant.uuid)
-          return "participant-api-key";
-        return undefined;
-      }
-    );
-  }
+  // if (mocks.userConnectionManager) {
+  // }
 }