@signe/room 2.10.0 → 3.0.0

package/src/server.ts

@@ -33,6 +33,47 @@ type CreateRoomOptions = {
   throttleStorage?: number;
 };
 
+type SessionData = {
+  publicId: string;
+  state?: any;
+  created?: number;
+  connected?: boolean;
+  disconnectedAt?: number;
+};
+
+type TransferData = {
+  privateId: string;
+  publicId: string;
+  restored: boolean;
+  created: number;
+};
+
+type StorageMetrics = {
+  loadMs: number;
+  loadStateKeys: number;
+  loadLegacyKeys: number;
+  persistFlushes: number;
+  persistWrites: number;
+  persistDeletes: number;
+  persistLastFlushMs: number;
+  sessionGcRuns: number;
+  sessionGcScanned: number;
+  sessionGcExpired: number;
+  sessionIndexRepairs: number;
+  transferGcRuns: number;
+  transferGcScanned: number;
+  transferGcExpired: number;
+};
+
+const STATE_PREFIX = "state:";
+const SESSION_PREFIX = "session:";
+const SESSION_PUBLIC_PREFIX = "session-public:";
+const TRANSFER_PREFIX = "transfer:";
+const INTERNAL_PREFIX = "$room:";
+const SESSION_GC_LAST_RUN_KEY = `${INTERNAL_PREFIX}session-gc:last-run`;
+const TRANSFER_GC_LAST_RUN_KEY = `${INTERNAL_PREFIX}transfer-gc:last-run`;
+const DEFAULT_TRANSFER_EXPIRY_MS = 5 * 60 * 1000;
+
 /**
  * @class Server
  * @implements {Party.Server}
@@ -89,6 +130,111 @@ export class Server implements Party.Server {
     return this.room.storage
   }
 
+  private stateKey(path: string) {
+    return `${STATE_PREFIX}${path}`;
+  }
+
+  private sessionKey(privateId: string) {
+    return `${SESSION_PREFIX}${privateId}`;
+  }
+
+  private sessionPublicKey(publicId: string) {
+    return `${SESSION_PUBLIC_PREFIX}${publicId}`;
+  }
+
+  private transferKey(token: string) {
+    return `${TRANSFER_PREFIX}${token}`;
+  }
+
+  private isInternalStorageKey(key: string) {
+    return key.startsWith(STATE_PREFIX)
+      || key.startsWith(SESSION_PREFIX)
+      || key.startsWith(SESSION_PUBLIC_PREFIX)
+      || key.startsWith(TRANSFER_PREFIX)
+      || key.startsWith(INTERNAL_PREFIX);
+  }
+
+  private async listStorage<T = unknown>(prefix?: string) {
+    if (!prefix) {
+      return this.room.storage.list<T>();
+    }
+    return this.room.storage.list<T>({ prefix });
+  }
+
+  private async loadStatePath<T = unknown>(path: string) {
+    return this.room.storage.get<T>(this.stateKey(path));
+  }
+
+  private async saveStatePath(path: string, value: any) {
+    await this.room.storage.put(this.stateKey(path), value);
+  }
+
+  private async putStorageEntries(entries: Record<string, any>) {
+    const keys = Object.keys(entries);
+    if (keys.length === 0) return;
+    if (keys.length === 1) {
+      const key = keys[0];
+      await this.room.storage.put(key, entries[key]);
+      return;
+    }
+    await this.room.storage.put(entries);
+  }
+
+  private async deleteStorageKeys(keys: string[]) {
+    const uniqueKeys = Array.from(new Set(keys));
+    if (uniqueKeys.length === 0) return;
+    if (uniqueKeys.length === 1) {
+      await this.room.storage.delete(uniqueKeys[0]);
+      return;
+    }
+    await this.room.storage.delete(uniqueKeys);
+  }
+
+  private async deleteStatePath(path: string) {
+    const stateKey = this.stateKey(path);
+    const descendantEntries = await this.listStorage(`${stateKey}.`);
+    await this.deleteStorageKeys([
+      ...Array.from(descendantEntries.keys()),
+      path,
+    ]);
+    await this.saveStatePath(path, DELETE_TOKEN);
+  }
+
+  private createStorageMetrics(): StorageMetrics {
+    return {
+      loadMs: 0,
+      loadStateKeys: 0,
+      loadLegacyKeys: 0,
+      persistFlushes: 0,
+      persistWrites: 0,
+      persistDeletes: 0,
+      persistLastFlushMs: 0,
+      sessionGcRuns: 0,
+      sessionGcScanned: 0,
+      sessionGcExpired: 0,
+      sessionIndexRepairs: 0,
+      transferGcRuns: 0,
+      transferGcScanned: 0,
+      transferGcExpired: 0,
+    };
+  }
+
+  private getTransferExpiryTime(subRoom: any) {
+    return subRoom?.transferExpiryTime
+      ?? subRoom?.constructor?.prototype?.transferExpiryTime
+      ?? subRoom?.constructor?.transferExpiryTime
+      ?? DEFAULT_TRANSFER_EXPIRY_MS;
+  }
+
+  private getPrivateId(conn: Party.Connection) {
+    return (conn.state as any)?.privateId || conn.sessionId || conn.id;
+  }
+
+  private hasActiveSessionConnection(privateId: string) {
+    return Array.from(this.room.getConnections())
+      .some((conn) => this.getPrivateId(conn) === privateId);
+  }
+
   async send(conn: Party.Connection, obj: any, subRoom: any) {
     obj = structuredClone(obj);
     if (subRoom.interceptorPacket) {
@@ -138,28 +284,37 @@ export class Server implements Party.Server {
     const subRoom = await this.getSubRoom();
     if (!subRoom) return;
 
+    const SESSION_EXPIRY_TIME = Number(options.sessionExpiryTime);
+    if (!Number.isFinite(SESSION_EXPIRY_TIME)) {
+      return;
+    }
+
     // Get active connections
     const activeConnections = [...this.room.getConnections()];
-    const activePrivateIds = new Set(activeConnections.map(conn => conn.id));
+    const activePrivateIds = new Set(activeConnections.map(conn => this.getPrivateId(conn)));
 
     try {
       // Get all sessions from storage
-      const sessions = await this.room.storage.list();
+      const sessions = await this.listStorage<SessionData>(SESSION_PREFIX);
       const users = this.getUsersProperty(subRoom);
       const usersPropName = this.getUsersPropName(subRoom);
+      const metrics = subRoom.$storageMetrics as StorageMetrics | undefined;
+      if (metrics) {
+        metrics.sessionGcRuns += 1;
+        metrics.sessionGcScanned += sessions.size;
+      }
 
       // Store valid publicIds from sessions
       const validPublicIds = new Set<string>();
       const expiredPublicIds = new Set<string>();
-      const SESSION_EXPIRY_TIME = options.sessionExpiryTime
       const now = Date.now();
 
       for (const [key, session] of sessions) {
         // Only process session entries
-        if (!key.startsWith('session:')) continue;
+        if (!key.startsWith(SESSION_PREFIX)) continue;
 
-        const privateId = key.replace('session:', '');
-        const typedSession = session as { publicId: string, created: number, connected: boolean };
+        const privateId = key.slice(SESSION_PREFIX.length);
+        const typedSession = session as SessionData;
 
         // Check if session should be deleted based on:
         // 1. Connection is not active
@@ -167,16 +322,22 @@ export class Server implements Party.Server {
         // 3. Session is older than expiry time
         if (!activePrivateIds.has(privateId) &&
           !typedSession.connected &&
-          (now - typedSession.created) > SESSION_EXPIRY_TIME) {
+          typedSession.disconnectedAt !== undefined &&
+          (now - typedSession.disconnectedAt) >= SESSION_EXPIRY_TIME) {
           // Delete expired session
           await this.deleteSession(privateId);
           expiredPublicIds.add(typedSession.publicId);
+          if (metrics) {
+            metrics.sessionGcExpired += 1;
+          }
         } else if (typedSession && typedSession.publicId) {
           // Keep track of valid publicIds from active or recent sessions
           validPublicIds.add(typedSession.publicId);
         }
       }
 
+      await this.repairSessionPublicIndexes(sessions, metrics);
+
       // Clean up users only if ALL their sessions are expired
       if (users && usersPropName) {
         const currentUsers = users();
@@ -184,7 +345,6 @@ export class Server implements Party.Server {
           // Only delete user if they have an expired session and no valid sessions
           if (expiredPublicIds.has(publicId) && !validPublicIds.has(publicId)) {
             delete currentUsers[publicId];
-            await this.room.storage.delete(`${usersPropName}.${publicId}`);
           }
         }
       }
@@ -194,6 +354,177 @@ export class Server implements Party.Server {
     }
   }
 
+  private scheduleSessionGarbageCollector(sessionExpiryTime: number | undefined, privateId?: string) {
+    const normalizedSessionExpiryTime = Number(sessionExpiryTime);
+    if (!Number.isFinite(normalizedSessionExpiryTime) || normalizedSessionExpiryTime < 0) {
+      return;
+    }
+
+    setTimeout(() => {
+      if (privateId) {
+        void this.expireDisconnectedSession(privateId, normalizedSessionExpiryTime);
+        return;
+      }
+      void this.garbageCollector({ sessionExpiryTime: normalizedSessionExpiryTime });
+    }, normalizedSessionExpiryTime);
+  }
+
+  private getSessionExpiryTime(subRoom: any) {
+    return subRoom?.sessionExpiryTime
+      ?? subRoom?.constructor?.prototype?.sessionExpiryTime
+      ?? subRoom?.constructor?.sessionExpiryTime;
+  }
+
+  private async shouldRunSessionGarbageCollector(sessionExpiryTime: number | undefined) {
+    const normalizedSessionExpiryTime = Number(sessionExpiryTime);
+    if (!Number.isFinite(normalizedSessionExpiryTime) || normalizedSessionExpiryTime < 0) {
+      return false;
+    }
+
+    const now = Date.now();
+    const lastRun = await this.room.storage.get<number>(SESSION_GC_LAST_RUN_KEY);
+    if (lastRun && now - lastRun < normalizedSessionExpiryTime) {
+      return false;
+    }
+
+    await this.room.storage.put(SESSION_GC_LAST_RUN_KEY, now);
+    return true;
+  }
+
+  private async shouldRunInterval(key: string, interval: number) {
+    const normalizedInterval = Number(interval);
+    if (!Number.isFinite(normalizedInterval) || normalizedInterval < 0) {
+      return false;
+    }
+
+    const now = Date.now();
+    const lastRun = await this.room.storage.get<number>(key);
+    if (lastRun && now - lastRun < normalizedInterval) {
+      return false;
+    }
+
+    await this.room.storage.put(key, now);
+    return true;
+  }
+
+  private async repairSessionPublicIndexes(
+    sessions?: Map<string, SessionData>,
+    metrics?: StorageMetrics
+  ) {
+    const sessionEntries = sessions ?? await this.listStorage<SessionData>(SESSION_PREFIX);
+    const expected = new Map<string, string[]>();
+
+    for (const [key, session] of sessionEntries) {
+      if (!session?.publicId) continue;
+      const privateId = key.slice(SESSION_PREFIX.length);
+      const privateIds = expected.get(session.publicId) ?? [];
+      privateIds.push(privateId);
+      expected.set(session.publicId, privateIds);
+    }
+
+    const publicIndexes = await this.listStorage<string[]>(SESSION_PUBLIC_PREFIX);
+    const writes: Record<string, string[]> = {};
+    const deletes: string[] = [];
+    let repairs = 0;
+
+    for (const [publicId, privateIds] of expected) {
+      privateIds.sort();
+      const key = this.sessionPublicKey(publicId);
+      const existing = publicIndexes.get(key) ?? [];
+      const normalizedExisting = Array.isArray(existing) ? [...existing].sort() : [];
+      if (JSON.stringify(normalizedExisting) !== JSON.stringify(privateIds)) {
+        writes[key] = privateIds;
+        repairs += 1;
+      }
+    }
+
+    for (const [key] of publicIndexes) {
+      const publicId = key.slice(SESSION_PUBLIC_PREFIX.length);
+      if (!expected.has(publicId)) {
+        deletes.push(key);
+        repairs += 1;
+      }
+    }
+
+    await Promise.all([
+      this.putStorageEntries(writes),
+      this.deleteStorageKeys(deletes),
+    ]);
+
+    if (metrics) {
+      metrics.sessionIndexRepairs += repairs;
+    }
+  }
+
+  private isTransferExpired(transfer: Partial<TransferData> | null | undefined, transferExpiryTime: number) {
+    if (!transfer) return true;
+    if (!Number.isFinite(transferExpiryTime) || transferExpiryTime < 0) {
+      return false;
+    }
+    const created = Number(transfer.created);
+    if (!Number.isFinite(created)) {
+      return false;
+    }
+    return Date.now() - created >= transferExpiryTime;
+  }
+
+  private async cleanupExpiredTransfers(transferExpiryTime: number, metrics?: StorageMetrics) {
+    if (!Number.isFinite(transferExpiryTime) || transferExpiryTime < 0) {
+      return;
+    }
+
+    const transfers = await this.listStorage<TransferData>(TRANSFER_PREFIX);
+    const expiredKeys: string[] = [];
+    for (const [key, transfer] of transfers) {
+      if (this.isTransferExpired(transfer, transferExpiryTime)) {
+        expiredKeys.push(key);
+      }
+    }
+
+    await this.deleteStorageKeys(expiredKeys);
+
+    if (metrics) {
+      metrics.transferGcRuns += 1;
+      metrics.transferGcScanned += transfers.size;
+      metrics.transferGcExpired += expiredKeys.length;
+    }
+  }
+
+  private async expireDisconnectedSession(privateId: string, sessionExpiryTime: number) {
+    const session = await this.getSession(privateId);
+    if (!session || session.connected || session.disconnectedAt === undefined) {
+      return;
+    }
+
+    if (this.hasActiveSessionConnection(privateId)) {
+      return;
+    }
+
+    const elapsed = Date.now() - session.disconnectedAt;
+    if (elapsed < sessionExpiryTime) {
+      setTimeout(() => {
+        void this.expireDisconnectedSession(privateId, sessionExpiryTime);
+      }, sessionExpiryTime - elapsed);
+      return;
+    }
+
+    await this.deleteSession(privateId);
+
+    const privateIds = await this.getSessionPrivateIds(session.publicId);
+    for (const otherPrivateId of privateIds) {
+      const otherSession = await this.getSession(otherPrivateId);
+      if (otherSession?.publicId === session.publicId) {
+        return;
+      }
+    }
+
+    const subRoom = await this.getSubRoom();
+    const users = this.getUsersProperty(subRoom);
+    if (users?.()[session.publicId]) {
+      delete users()[session.publicId];
+    }
+  }
+
   /**
    * @method createRoom
    * @private
@@ -231,22 +562,61 @@ export class Server implements Party.Server {
     // Load the room's memory from storage
     // This ensures persistence across server restarts
     const loadMemory = async () => {
-      const root = await this.room.storage.get(".");
-      const memory = await this.room.storage.list();
+      const startedAt = Date.now();
+      const metrics = instance.$storageMetrics as StorageMetrics;
+      const root = await this.loadStatePath(".");
+      const memory = await this.listStorage(STATE_PREFIX);
+      metrics.loadStateKeys = memory.size;
       const tmpObject: any = root || {};
-      for (let [key, value] of memory) {
-        if (key.startsWith('session:')) {
-          continue;
-        }
-        if (key == ".") {
+      for (let [storageKey, value] of memory) {
+        const key = storageKey.slice(STATE_PREFIX.length);
+        if (key === ".") {
           continue;
         }
         dset(tmpObject, key, value);
       }
+
+      if (root === undefined && memory.size === 0) {
+        const legacyRoot = await this.room.storage.get(".");
+        const legacyMemory = await this.room.storage.list();
+        const legacyObject: any = legacyRoot || {};
+        const migratedEntries: Array<[string, any]> = [];
+        const legacyDeleteKeys: string[] = [];
+
+        if (legacyRoot !== undefined) {
+          migratedEntries.push([".", legacyRoot]);
+          legacyDeleteKeys.push(".");
+        }
+
+        for (let [key, value] of legacyMemory) {
+          if (key === "." || this.isInternalStorageKey(key)) {
+            continue;
+          }
+          dset(legacyObject, key, value);
+          migratedEntries.push([key, value]);
+          legacyDeleteKeys.push(key);
+        }
+
+        metrics.loadLegacyKeys = migratedEntries.length;
+        await this.putStorageEntries(
+          Object.fromEntries(
+            migratedEntries.map(([path, value]) => [this.stateKey(path), value])
+          )
+        );
+        if (legacyRoot !== undefined) {
+          await this.deleteStorageKeys(legacyDeleteKeys);
+        }
+        load(instance, legacyObject, true);
+        metrics.loadMs = Date.now() - startedAt;
+        return;
+      }
+
       load(instance, tmpObject, true);
+      metrics.loadMs = Date.now() - startedAt;
     };
 
     instance.$memoryAll = {}
+    instance.$storageMetrics = this.createStorageMetrics();
     instance.$autoSync = instance["autoSync"] !== false; // Default to true
     instance.$pendingSync = new Map<string, any>();
     instance.$pendingInitialSync = new Map<Party.Connection, string>(); // Store connections waiting for initial sync with their publicId
@@ -333,17 +703,9 @@ export class Server implements Party.Server {
         return null;
       }
 
-      const sessions = await this.room.storage.list();
-      let userSession: any = null;
-      let privateId: string | null = null;
-
-      for (const [key, session] of sessions) {
-        if (key.startsWith('session:') && (session as any).publicId === publicId) {
-          userSession = session;
-          privateId = key.replace('session:', '');
-          break;
-        }
-      }
+      const sessionEntry = await this.getSessionEntryByPublicId(publicId);
+      const userSession = sessionEntry?.session;
+      const privateId = sessionEntry?.privateId ?? null;
 
       if (!userSession || !privateId) {
         console.error(`[sessionTransfer] Session for publicId ${publicId} not found.`);
@@ -428,23 +790,91 @@ export class Server implements Party.Server {
         values.clear();
         return;
       }
+      const startedAt = Date.now();
+      const stateWrites: Record<string, any> = {};
+      const deleteTasks: Promise<void>[] = [];
       for (let [path, value] of values) {
         const _instance =
           path == "." ? instance : getByPath(instance, path);
-        const itemValue = createStatesSnapshot(_instance);
         if (value == DELETE_TOKEN) {
-          await this.room.storage.delete(path);
+          deleteTasks.push(this.deleteStatePath(path));
         } else {
-          await this.room.storage.put(path, itemValue);
+          const itemValue = _instance?.$snapshot
+            ? createStatesSnapshot(_instance)
+            : value;
+          stateWrites[this.stateKey(path)] = itemValue;
         }
       }
+      await Promise.all([
+        this.putStorageEntries(stateWrites),
+        ...deleteTasks,
+      ]);
+      const metrics = instance.$storageMetrics as StorageMetrics | undefined;
+      if (metrics) {
+        metrics.persistFlushes += 1;
+        metrics.persistWrites += Object.keys(stateWrites).length;
+        metrics.persistDeletes += deleteTasks.length;
+        metrics.persistLastFlushMs = Date.now() - startedAt;
+      }
       values.clear();
     }
 
+    const debouncePersist = (wait: number) => {
+      let timeout: ReturnType<typeof setTimeout> | null = null;
+      let flushing = false;
+      const pending = new Map<string, any>();
+
+      const schedule = () => {
+        if (timeout) {
+          clearTimeout(timeout);
+        }
+        timeout = setTimeout(() => {
+          void flush();
+        }, wait);
+      };
+
+      const flush = async () => {
+        timeout = null;
+        if (flushing) {
+          schedule();
+          return;
+        }
+
+        const values = new Map(pending);
+        pending.clear();
+        if (!values.size) {
+          return;
+        }
+
+        flushing = true;
+        try {
+          await persistCb(values);
+        } finally {
+          flushing = false;
+          if (pending.size) {
+            schedule();
+          }
+        }
+      };
+
+      return (values: Map<string, any>) => {
+        if (initPersist) {
+          values.clear();
+          return;
+        }
+
+        for (const [path, value] of values) {
+          pending.set(path, value);
+        }
+        values.clear();
+        schedule();
+      };
+    };
+
     // Set up syncing and persistence with throttling to optimize performance
     syncClass(instance, {
       onSync: instance["throttleSync"] ? throttle(syncCb, instance["throttleSync"]) : syncCb,
-      onPersist: instance["throttleStorage"] ? throttle(persistCb, instance["throttleStorage"]) : persistCb,
+      onPersist: instance["throttleStorage"] ? debouncePersist(instance["throttleStorage"]) : persistCb,
     });
 
     await loadMemory();
@@ -564,29 +994,107 @@ export class Server implements Party.Server {
    * console.log(session);
    * ```
    */
-  async getSession(privateId: string): Promise<{ publicId: string, state?: any, created?: number, connected?: boolean } | null> {
+  async getSession(privateId: string): Promise<SessionData | null> {
     if (!privateId) return null;
     try {
-      const session = await this.room.storage.get(`session:${privateId}`);
-      return session as { publicId: string, state?: any, created: number, connected: boolean } | null;
+      const session = await this.room.storage.get(this.sessionKey(privateId));
+      return session as SessionData | null;
     } catch (e) {
       return null;
     }
   }
 
-  private async saveSession(privateId: string, data: { publicId: string, state?: any, created?: number, connected?: boolean }) {
+  private async getSessionPrivateIds(publicId: string): Promise<string[]> {
+    if (!publicId) return [];
+    const privateIds = await this.room.storage.get<string[]>(this.sessionPublicKey(publicId));
+    return Array.isArray(privateIds) ? privateIds : [];
+  }
+
+  private async saveSessionPrivateIds(publicId: string, privateIds: string[]) {
+    const key = this.sessionPublicKey(publicId);
+    if (privateIds.length === 0) {
+      await this.room.storage.delete(key);
+      return;
+    }
+    await this.room.storage.put(key, privateIds);
+  }
+
+  private async addSessionToPublicIndex(privateId: string, publicId: string) {
+    const privateIds = await this.getSessionPrivateIds(publicId);
+    if (privateIds.includes(privateId)) {
+      return;
+    }
+    await this.saveSessionPrivateIds(publicId, [...privateIds, privateId]);
+  }
+
+  private async removeSessionFromPublicIndex(privateId: string, publicId: string) {
+    const privateIds = await this.getSessionPrivateIds(publicId);
+    await this.saveSessionPrivateIds(
+      publicId,
+      privateIds.filter((id) => id !== privateId)
+    );
+  }
+
+  private async getSessionEntryByPublicId(publicId: string): Promise<{ privateId: string; session: SessionData } | null> {
+    const indexedPrivateIds = await this.getSessionPrivateIds(publicId);
+    const stalePrivateIds: string[] = [];
+
+    for (const privateId of indexedPrivateIds) {
+      const session = await this.getSession(privateId);
+      if (session?.publicId === publicId) {
+        return { privateId, session };
+      }
+      stalePrivateIds.push(privateId);
+    }
+
+    if (stalePrivateIds.length) {
+      await this.saveSessionPrivateIds(
+        publicId,
+        indexedPrivateIds.filter((id) => !stalePrivateIds.includes(id))
+      );
+    }
+
+    const sessions = await this.listStorage<SessionData>(SESSION_PREFIX);
+    for (const [key, session] of sessions) {
+      const privateId = key.slice(SESSION_PREFIX.length);
+      if (session?.publicId) {
+        await this.addSessionToPublicIndex(privateId, session.publicId);
+      }
+      if (session?.publicId === publicId) {
+        return { privateId, session };
+      }
+    }
+
+    return null;
+  }
+
+  private async saveSession(privateId: string, data: SessionData) {
+    const existingSession = await this.getSession(privateId);
     const sessionData = {
       ...data,
       created: data.created || Date.now(),
       connected: data.connected !== undefined ? data.connected : true
     };
-    await this.room.storage.put(`session:${privateId}`, sessionData);
+    await this.room.storage.put(this.sessionKey(privateId), sessionData);
+    if (existingSession?.publicId && existingSession.publicId !== sessionData.publicId) {
+      await this.removeSessionFromPublicIndex(privateId, existingSession.publicId);
+    }
+    await this.addSessionToPublicIndex(privateId, sessionData.publicId);
   }
 
   private async updateSessionConnection(privateId: string, connected: boolean) {
     const session = await this.getSession(privateId);
     if (session) {
-      await this.saveSession(privateId, { ...session, connected });
+      const nextSession = { ...session, connected };
+      if (connected) {
+        delete nextSession.disconnectedAt;
+      } else {
+        nextSession.disconnectedAt = Date.now();
+      }
+      if (!await this.getSession(privateId)) {
+        return;
+      }
+      await this.saveSession(privateId, nextSession);
     }
   }
 
@@ -602,7 +1110,11 @@ export class Server implements Party.Server {
    * ```
    */
   async deleteSession(privateId: string) {
-    await this.room.storage.delete(`session:${privateId}`);
+    const session = await this.getSession(privateId);
+    await this.room.storage.delete(this.sessionKey(privateId));
+    if (session?.publicId) {
+      await this.removeSessionFromPublicIndex(privateId, session.publicId);
+    }
   }
 
   async onConnectClient(conn: Party.Connection, ctx: Party.ConnectionContext) {
@@ -615,8 +1127,17 @@ export class Server implements Party.Server {
       return;
     }
 
-    const sessionExpiryTime = subRoom.constructor.sessionExpiryTime;
-    await this.garbageCollector({ sessionExpiryTime });
+    const sessionExpiryTime = this.getSessionExpiryTime(subRoom);
+    if (await this.shouldRunSessionGarbageCollector(sessionExpiryTime)) {
+      await this.garbageCollector({ sessionExpiryTime });
+    }
+    const transferExpiryTime = this.getTransferExpiryTime(subRoom);
+    if (await this.shouldRunInterval(TRANSFER_GC_LAST_RUN_KEY, transferExpiryTime)) {
+      await this.cleanupExpiredTransfers(
+        transferExpiryTime,
+        subRoom.$storageMetrics as StorageMetrics | undefined
+      );
+    }
 
     // Check room guards
     const roomGuards = subRoom.constructor['_roomGuards'] || [];
@@ -636,14 +1157,20 @@ export class Server implements Party.Server {
     }
     let transferData: any = null;
     if (transferToken) {
-      transferData = await this.room.storage.get(`transfer:${transferToken}`);
+      const transferKey = this.transferKey(transferToken);
+      transferData = await this.room.storage.get<TransferData>(transferKey);
       if (transferData) {
-        await this.room.storage.delete(`transfer:${transferToken}`);
+        if (this.isTransferExpired(transferData, transferExpiryTime)) {
+          transferData = null;
+        }
+        await this.room.storage.delete(transferKey);
       }
     }
 
     // Check for existing session
-    const existingSession = await this.getSession(conn.id)
+    const requestedPrivateId = this.getPrivateId(conn);
+    const privateId = transferData?.privateId || requestedPrivateId;
+    const existingSession = await this.getSession(privateId)
 
     // Generate IDs
     const publicId = existingSession?.publicId || transferData?.publicId || generateShortUUID();
@@ -664,7 +1191,7 @@ export class Server implements Party.Server {
           user = isClass(classType) ? new classType() : classType(conn, ctx);
           signal()[publicId] = user;
           const snapshot = createStatesSnapshotDeep(user);
-          this.room.storage.put(`${usersPropName}.${publicId}`, snapshot);
+          await this.saveStatePath(`${usersPropName}.${publicId}`, snapshot);
         }
       }
       else {
@@ -674,13 +1201,12 @@ export class Server implements Party.Server {
       // Only store new session if it doesn't exist
       if (!existingSession) {
         // Use the transferred privateId if available, otherwise use connection id
-        const sessionPrivateId = transferData?.privateId || conn.id;
-        await this.saveSession(sessionPrivateId, {
+        await this.saveSession(privateId, {
           publicId
         });
       }
       else {
-        await this.updateSessionConnection(conn.id, true);
+        await this.updateSessionConnection(privateId, true);
       }
     }
     // Update user connection status if applicable
@@ -689,7 +1215,8 @@ export class Server implements Party.Server {
      // Store both IDs in connection state
      conn.setState({
       ...conn.state,
-      publicId
+      publicId,
+      privateId
     });
 
     // Call the room's onJoin method if it exists
@@ -729,6 +1256,10 @@ export class Server implements Party.Server {
    */
   async onConnect(conn: Party.Connection, ctx: Party.ConnectionContext) {
     if (ctx.request?.headers.has('x-shard-id')) {
+      if (!this.isAuthorizedShardRequest(ctx.request)) {
+        conn.close();
+        return;
+      }
       this.onConnectShard(conn, ctx);
     }
     else {
@@ -736,6 +1267,13 @@ export class Server implements Party.Server {
     }
   }
 
+  private isAuthorizedShardRequest(req?: Party.Request) {
+    const shardSecret = this.room.env.SHARD_SECRET;
+    return typeof shardSecret === 'string'
+      && shardSecret.length > 0
+      && req?.headers.get('x-access-shard') === shardSecret;
+  }
+
   /**
    * @method onConnectShard
    * @private
@@ -747,9 +1285,11 @@ export class Server implements Party.Server {
   onConnectShard(conn: Party.Connection, ctx: Party.ConnectionContext) {
     // Set shard metadata in connection state
     const shardId = ctx.request?.headers.get('x-shard-id') || 'unknown-shard';
+    const worldId = ctx.request?.headers.get('x-shard-world-id') || 'world-default';
     conn.setState({
       shard: true,
       shardId,
+      worldId,
       clients: new Map() // Track clients connected through this shard
     });
   }
@@ -1085,14 +1625,19 @@ export class Server implements Party.Server {
       return;
     }
 
-    const privateId = conn.id;
+    const privateId = this.getPrivateId(conn);
     const { publicId } = conn.state as any;
     const user = signal?.()[publicId];
 
     if (!user) return;
 
+    if (this.hasActiveSessionConnection(privateId)) {
+      return;
+    }
+
     // Mark session as disconnected instead of deleting it
     await this.updateSessionConnection(privateId, false);
+    this.scheduleSessionGarbageCollector(this.getSessionExpiryTime(subRoom), privateId);
 
     // Update user connection status in the signal
     const connectionUpdated = this.updateUserConnectionStatus(user, false);
@@ -1142,6 +1687,9 @@ export class Server implements Party.Server {
     }
 
     if (isFromShard) {
+      if (!this.isAuthorizedShardRequest(req)) {
+        return res.unauthorized('Invalid shard credentials');
+      }
       return this.handleShardRequest(req, res, shardId);
     }
 
@@ -1216,16 +1764,17 @@ export class Server implements Party.Server {
           load(user, hydratedSnapshot, true);
           
           // Save user snapshot to storage
-          await this.room.storage.put(`${usersPropName}.${publicId}`, userSnapshot);
+          await this.saveStatePath(`${usersPropName}.${publicId}`, userSnapshot);
         }
       }
 
       // Generate transfer token for the client to use when connecting
       const transferToken = generateShortUUID();
-      await this.room.storage.put(`transfer:${transferToken}`, {
+      await this.room.storage.put(this.transferKey(transferToken), {
         privateId,
         publicId,
-        restored: true
+        restored: true,
+        created: Date.now(),
       });
 
       return res.success({ transferToken });