@signe/room 2.10.0 → 3.0.0

package/src/session.guard.ts

@@ -1,5 +1,9 @@
 import type * as Party from "./types/party";
 
+function getPrivateId(sender: Party.Connection) {
+    return sender.sessionId || sender.id;
+}
+
 /**
  * @description Factory function that creates a session guard with access to room storage
  * @param {Party.Storage} storage - The room storage instance
@@ -29,7 +33,7 @@ export function createRequireSessionGuard(storage: Party.Storage) {
 
         try {
             // Check if session exists in storage
-            const session = await storage.get(`session:${sender.id}`);
+            const session = await storage.get(`session:${getPrivateId(sender)}`);
             
             // Return false if no session found
             if (!session) {
@@ -88,7 +92,7 @@ export const requireSession = async (sender: Party.Connection, value: any, room:
 
     try {
         // Check if session exists in storage
-        const session = await room.storage.get(`session:${sender.id}`);
+        const session = await room.storage.get(`session:${getPrivateId(sender)}`);
         
         // Return false if no session found
         if (!session) {

package/src/shard.ts

@@ -15,18 +15,44 @@ export interface ShardOptions {
 
 export class Shard {
   ws: PartyWebSocket;
-  connectionMap = new Map<string, Party.Connection>(); // Map privateId -> connection
+  connectionMap = new Map<string, Set<Party.Connection>>(); // Map privateId -> active connections
   mainServerStub: any;
   worldUrl: string | null = null;
-  worldId: string = 'default';
+  worldId: string;
   lastReportedConnections: number = 0;
   statsInterval: number = 30000; 
   statsIntervalId: any = null;
 
-  constructor(private room: Party.Room) {}
+  constructor(private room: Party.Room, options: ShardOptions = {}) {
+    this.worldUrl = options.worldUrl ?? null;
+    this.worldId = options.worldId
+      ?? this.getWorldIdFromShardId(room.id)
+      ?? this.getEnvString('WORLD_ID')
+      ?? this.getEnvString('SIGNE_WORLD_ID')
+      ?? 'world-default';
+    this.statsInterval = options.statsInterval ?? this.statsInterval;
+  }
+
+  private getPrivateId(conn: Party.Connection) {
+    return conn.sessionId || conn.id;
+  }
+
+  private getEnvString(key: string) {
+    const value = this.room.env?.[key];
+    return typeof value === 'string' && value.length > 0 ? value : undefined;
+  }
+
+  private getRoomIdFromShardId(shardId: string) {
+    return shardId.split(':')[0];
+  }
+
+  private getWorldIdFromShardId(shardId: string) {
+    const parts = shardId.split(':');
+    return parts.length >= 3 ? parts[1] : undefined;
+  }
 
   async onStart() {
-    const roomId = this.room.id.split(':')[0];
+    const roomId = this.getRoomIdFromShardId(this.room.id);
     const roomStub = this.room.context.parties.main.get(roomId);
     if (!roomStub) {
       console.warn('No room room stub found in main party context');
@@ -36,7 +62,9 @@ export class Shard {
     this.mainServerStub = roomStub;
     this.ws = await roomStub.socket({
       headers: {
-        'x-shard-id': this.room.id
+        'x-shard-id': this.room.id,
+        'x-shard-world-id': this.worldId,
+        'x-access-shard': this.room.env.SHARD_SECRET as string
       }
     }) as unknown as PartyWebSocket;
     
@@ -45,13 +73,25 @@ export class Shard {
       try {
         const message = JSON.parse(event.data);
 
+        if (message.type === 'shard.closeClient' && message.privateId) {
+          const clientConnections = this.connectionMap.get(message.privateId);
+          if (clientConnections?.size) {
+            for (const clientConn of [...clientConnections]) {
+              clientConn.close();
+            }
+          }
+          return;
+        }
+
         // If the message is directed to a specific client, forward it
         if (message.targetClientId) {
-          const clientConn = this.connectionMap.get(message.targetClientId);
-          if (clientConn) {
+          const clientConnections = this.connectionMap.get(message.targetClientId);
+          if (clientConnections?.size) {
             // Remove the routing information before forwarding
             delete message.targetClientId;
-            clientConn.send(message.data);
+            for (const clientConn of clientConnections) {
+              clientConn.send(message.data);
+            }
           }
         } else {
           // Broadcast to all clients if no specific target
@@ -62,12 +102,12 @@ export class Shard {
       }
     });
 
-    await this.updateWorldStats();
+    await this.updateWorldStats(true);
     this.startPeriodicStatsUpdates();
   }
 
   private startPeriodicStatsUpdates() {
-    if (!this.worldUrl) {
+    if (this.statsInterval <= 0 || !this.room.context.parties.world) {
       return;
     }
 
@@ -76,10 +116,11 @@ export class Shard {
     }
 
     this.statsIntervalId = setInterval(() => {
-      this.updateWorldStats().catch(error => {
+      this.updateWorldStats(true).catch(error => {
         console.error('Error in periodic stats update:', error);
       });
     }, this.statsInterval);
+    this.statsIntervalId?.unref?.();
   }
 
   private stopPeriodicStatsUpdates() {
@@ -90,8 +131,12 @@ export class Shard {
   }
 
   onConnect(conn: Party.Connection, ctx: Party.ConnectionContext) {
+    const privateId = this.getPrivateId(conn);
+
     // Store connection mapping
-    this.connectionMap.set(conn.id, conn);
+    const connections = this.connectionMap.get(privateId) ?? new Set<Party.Connection>();
+    connections.add(conn);
+    this.connectionMap.set(privateId, connections);
     
     // Capture all headers and request information
     const headers: Record<string, string> = {};
@@ -111,7 +156,7 @@ export class Shard {
     // Notify the main server about the new connection with complete connection metadata
     this.ws.send(JSON.stringify({
       type: 'shard.clientConnected',
-      privateId: conn.id,
+      privateId,
       requestInfo
     }));
 
@@ -126,7 +171,7 @@ export class Shard {
       // Wrap the original message with sender information
       const wrappedMessage = JSON.stringify({
         type: 'shard.clientMessage',
-        privateId: sender.id,
+        privateId: this.getPrivateId(sender),
         publicId: (sender.state as any)?.publicId,
         payload: parsedMessage
       });
@@ -139,28 +184,48 @@ export class Shard {
   }
 
   onClose(conn: Party.Connection) {
+    const privateId = this.getPrivateId(conn);
+
     // Remove connection from the map
-    this.connectionMap.delete(conn.id);
+    const connections = this.connectionMap.get(privateId);
+    connections?.delete(conn);
+
+    if (connections?.size) {
+      this.updateWorldStats();
+      return;
+    }
+
+    this.connectionMap.delete(privateId);
     
     // Notify main server about disconnection
     this.ws.send(JSON.stringify({
       type: 'shard.clientDisconnected',
-      privateId: conn.id,
+      privateId,
       publicId: (conn.state as any)?.publicId
     }));
 
     this.updateWorldStats();
   }
 
-  async updateWorldStats(): Promise<boolean> {
-    const currentConnections = this.connectionMap.size;
+  async updateWorldStats(force = false): Promise<boolean> {
+    const currentConnections = Array.from(this.connectionMap.values())
+      .reduce((total, connections) => total + connections.size, 0);
     
-    if (currentConnections === this.lastReportedConnections) {
+    if (!force && currentConnections === this.lastReportedConnections) {
       return true;
     }
 
     try {
-      const worldRoom = this.room.context.parties.world.get('world-default');
+      const worldParty = this.room.context.parties.world;
+      if (!worldParty) {
+        return false;
+      }
+
+      const worldRoom = worldParty.get(this.worldId);
+      if (!worldRoom?.fetch) {
+        return false;
+      }
+
       const response = await worldRoom.fetch('/update-shard', {
         method: 'POST',
         headers: {
@@ -169,6 +234,7 @@ export class Shard {
         },
         body: JSON.stringify({
           shardId: this.room.id,
+          worldId: this.worldId,
           connections: currentConnections
         })
       });
@@ -219,7 +285,9 @@ export class Shard {
       
       // Add shard identification
       headers.set('x-shard-id', this.room.id);
+      headers.set('x-shard-world-id', this.worldId);
       headers.set('x-forwarded-by-shard', 'true');
+      headers.set('x-access-shard', this.room.env.SHARD_SECRET as string);
       
       // Client IP tracking for the main server
       const clientIp = req.headers.get('x-forwarded-for') || 'unknown';
@@ -234,7 +302,7 @@ export class Shard {
         body
       };
       // Forward the request to the main server
-      const response = await this.mainServerStub.fetch(path, requestInit);
+      const response = await this.mainServerStub.fetch(path + url.search, requestInit);
       return response;
     } catch (error) {
       return response(500, { error: 'Error forwarding request' });
@@ -247,8 +315,8 @@ export class Shard {
    * @description Executed periodically, used to perform maintenance tasks
    */
   async onAlarm() {
-    await this.updateWorldStats();
+    await this.updateWorldStats(true);
   }
 }
 
-Shard satisfies Party.Worker;
+Shard satisfies Party.Worker;

package/src/storage.ts

@@ -1,15 +1,39 @@
 export class Storage {
   private memory = new Map();
-  async put(key, value) {
-    this.memory.set(key, value);
+  async put(key, value?) {
+    if (typeof key === "string") {
+      this.memory.set(key, value);
+      return;
+    }
+
+    for (const [entryKey, entryValue] of Object.entries(key)) {
+      this.memory.set(entryKey, entryValue);
+    }
   }
   async get(key) {
     return this.memory.get(key);
   }
   async delete(key) {
-    this.memory.delete(key);
+    if (Array.isArray(key)) {
+      let deleted = 0;
+      for (const entryKey of key) {
+        if (this.memory.delete(entryKey)) {
+          deleted += 1;
+        }
+      }
+      return deleted;
+    }
+    return this.memory.delete(key);
   }
-  async list() {
-    return this.memory;
+  async list(options?: { prefix?: string }) {
+    if (!options?.prefix) {
+      return this.memory;
+    }
+
+    return new Map(
+      Array.from(this.memory.entries()).filter(([key]) =>
+        String(key).startsWith(options.prefix!)
+      )
+    );
   }
 }

package/src/testing.ts

@@ -30,6 +30,7 @@ import { Shard } from "./shard"
 export async function testRoom(Room, options: {
     hibernate?: boolean,
     shard?: boolean,
+    id?: string,
     env?: Record<string, string>,
     parties?: Record<string, (io: any) => any>,
     partyFn?: (io: any) => any
@@ -42,7 +43,7 @@ export async function testRoom(Room, options: {
     }
 
     const isShard = options.shard || false
-    const io = new ServerIo(Room.path, isShard ? {
+    const io = new ServerIo(options.id ?? Room.path, isShard ? {
         parties: {
             game: createServer,
             ...(options.parties || {})
@@ -92,7 +93,7 @@ export async function testRoom(Room, options: {
             return client
         },
         getServerUser: async (client: any, prop = 'users') => {
-            const privateId = client.conn.id;
+            const privateId = client.conn.sessionId || client.conn.id;
             const session = await (server as Server).getSession(privateId);
             return (server as any).subRoom[prop]()[session?.publicId];
         }
@@ -114,4 +115,4 @@ export async function request(room: Server | Shard, path: string, options: {
 
 export function tick(ms: number = 0) {
     return new Promise(resolve => setTimeout(resolve, ms))
-}
+}

package/src/types/party.ts

@@ -150,6 +150,9 @@ import type {
   export type Connection<TState = unknown> = WebSocket & {
     /** Connection identifier */
     id: string;
+
+    /** Stable private session identifier shared by reconnects or multiple tabs. */
+    sessionId?: string;
   
     /** @deprecated You can access the socket properties directly on the connection*/
     socket: WebSocket;
@@ -548,4 +551,4 @@ import type {
   export type PartyKitConnection = Connection;
   
   /** @deprecated Use `Party.ServerOptions` instead */
-  export type PartyServerOptions = ServerOptions;
+  export type PartyServerOptions = ServerOptions;

package/src/world.guard.ts

@@ -1,8 +1,7 @@
 import * as Party from "./types/party";
 import { JWTAuth } from "./jwt";
-import { response } from "./utils";
 
-export const guardManageWorld = async (_, req: Party.Request, room: Party.Room): Promise<boolean> => {
+export const guardManageWorld = async (_: unknown, req: Party.Request, room: Party.Room): Promise<boolean> => {
     const tokenShard = req.headers.get("x-access-shard");
     if (tokenShard) {
         if (tokenShard !== room.env.SHARD_SECRET) {
@@ -11,7 +10,7 @@ export const guardManageWorld = async (_, req: Party.Request, room: Party.Room):
         return true
     }
     const url = new URL(req.url);
-    const token = req.headers.get("Authorization") ?? url.searchParams.get("world-auth-token");
+    const token = getAuthToken(req, url);
     if (!token) {
         return false;
     }
@@ -21,8 +20,28 @@ export const guardManageWorld = async (_, req: Party.Request, room: Party.Room):
         if (!payload) {
             return false;
         }
+        if (!canAccessWorld(payload, room.id)) {
+            return false;
+        }
     } catch (error) {
         return false;
     }
     return true;
-}
+}
+
+function getAuthToken(req: Party.Request, url: URL) {
+    const authorization = req.headers.get("Authorization");
+    if (authorization?.startsWith("Bearer ")) {
+        return authorization.slice("Bearer ".length).trim();
+    }
+    return authorization ?? url.searchParams.get("world-auth-token");
+}
+
+function canAccessWorld(payload: Record<string, unknown>, worldId: string) {
+    const worlds = payload.worlds;
+    if (!Array.isArray(worlds)) {
+        return false;
+    }
+
+    return worlds.some((world) => world === "*" || world === worldId);
+}

package/src/world.ts

@@ -27,11 +27,14 @@ const RoomConfigSchema = z.object({
 const RegisterShardSchema = z.object({
   shardId: z.string(),
   roomId: z.string(),
+  worldId: z.string().optional(),
   url: z.string().url(),
   maxConnections: z.number().int().positive(),
 });
 
 const UpdateShardStatsSchema = z.object({
+  shardId: z.string(),
+  worldId: z.string().optional(),
   connections: z.number().int().min(0),
   status: z.enum(['active', 'maintenance', 'draining']).optional(),
 });
@@ -59,6 +62,7 @@ class RoomConfig {
 class ShardInfo {
   @id() id: string;
   @sync() roomId = signal("");
+  @sync() worldId = signal("");
   @sync() url = signal("");
   @sync({
     persist: false
@@ -95,8 +99,7 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
     if (!SHARD_SECRET) {
       throw new Error("SHARD_SECRET env variable is not set");
     }
-    // TODO
-    //setTimeout(() => this.cleanupInactiveShards(), 60000);
+    this.scheduleInactiveShardCleanup();
   }
 
   async onJoin(user: any, conn: Party.Connection, ctx: Party.ConnectionContext) {
@@ -115,6 +118,15 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
   }
 
   // Helper methods
+  private getWorldId() {
+    return this.room.id;
+  }
+
+  private scheduleInactiveShardCleanup() {
+    const timeoutId = setTimeout(() => this.cleanupInactiveShards(), 60000);
+    timeoutId?.unref?.();
+  }
+
   private cleanupInactiveShards() {
     const now = Date.now();
     const timeout = 5 * 60 * 1000; // 5 minutes timeout
@@ -130,7 +142,15 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
     });
 
     // Schedule next cleanup
-    setTimeout(() => this.cleanupInactiveShards(), 60000);
+    this.scheduleInactiveShardCleanup();
+  }
+
+  private removeShard(shardId: string) {
+    delete this.shards()[shardId];
+  }
+
+  private shouldCompleteDrain(shard: ShardInfo) {
+    return shard.status() === 'draining' && shard.currentConnections() === 0;
   }
 
   // Actions
@@ -139,8 +159,19 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
     method: 'POST',
   })
   @Guard([guardManageWorld])
-  async registerRoom(req: Party.Request) {
-    const roomConfig: z.infer<typeof RoomConfigSchema> = await req.json();
+  async registerRoom(req: Party.Request, res?: ServerResponse) {
+    const parseResult = RoomConfigSchema.safeParse(await req.json());
+    if (!parseResult.success) {
+      return res?.badRequest("Invalid room configuration", {
+        details: parseResult.error
+      });
+    }
+
+    const roomConfig = parseResult.data;
+    if (roomConfig.maxShards !== undefined && roomConfig.minShards > roomConfig.maxShards) {
+      return res?.badRequest("minShards cannot be greater than maxShards");
+    }
+
     const roomId = roomConfig.name;
 
     if (!this.rooms()[roomId]) {
@@ -170,6 +201,8 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
       room.minShards.set(roomConfig.minShards);
       room.maxShards.set(roomConfig.maxShards);
     }
+
+    await this.ensureMinShards(roomId);
   }
 
   @Request({
@@ -178,7 +211,14 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
   })
   @Guard([guardManageWorld])
   async updateShardStats(req: Party.Request, res: ServerResponse) {
-    const body: { shardId: string; connections: number; status: ShardStatus } = await req.json();
+    const parseResult = UpdateShardStatsSchema.safeParse(await req.json());
+    if (!parseResult.success) {
+      return res.badRequest("Invalid shard stats", {
+        details: parseResult.error
+      });
+    }
+
+    const body = parseResult.data;
     const { shardId, connections, status } = body;
     const shard = this.shards()[shardId];
 
@@ -186,11 +226,19 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
       return res.notFound(`Shard ${shardId} not found`);
     }
 
+    if (body.worldId && body.worldId !== this.getWorldId()) {
+      return res.badRequest(`Shard ${shardId} belongs to world ${body.worldId}, not ${this.getWorldId()}`);
+    }
+
     shard.currentConnections.set(connections);
     if (status) {
       shard.status.set(status);
     }
     shard.lastHeartbeat.set(Date.now());
+
+    if (this.shouldCompleteDrain(shard)) {
+      this.removeShard(shard.id);
+    }
   }
 
   @Request({
@@ -199,7 +247,14 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
   })
   @Guard([guardManageWorld])
   async scaleRoom(req: Party.Request, res: ServerResponse) {
-    const data: z.infer<typeof ScaleRoomSchema> = await req.json();
+    const parseResult = ScaleRoomSchema.safeParse(await req.json());
+    if (!parseResult.success) {
+      return res.badRequest("Invalid scale room request", {
+        details: parseResult.error
+      });
+    }
+
+    const data = parseResult.data;
     const { targetShardCount, shardTemplate, roomId } = data;
 
     // Validate room exists
@@ -223,8 +278,8 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
 
     // Handle scaling down
     if (targetShardCount < previousShardCount) {
-      // Find candidates for removal (prioritize draining or low-connection shards)
-      const shardsToRemove = [...roomShards]
+      // Find drain candidates (prioritize already-draining or low-connection shards)
+      const shardsToDrain = [...roomShards]
         .sort((a, b) => {
           // Prioritize draining status
           if (a.status() === 'draining' && b.status() !== 'draining') return -1;
@@ -235,14 +290,12 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
         })
         .slice(0, previousShardCount - targetShardCount);
 
-      // Remove the selected shards
-      const shardsToKeep = roomShards.filter(
-        shard => !shardsToRemove.some(s => s.id === shard.id)
-      );
-
-      // Update shards
-      for (const shard of shardsToRemove) {
-        delete this.shards()[shard.id];
+      // Complete empty drains immediately, otherwise stop routing new clients there.
+      for (const shard of shardsToDrain) {
+        shard.status.set('draining');
+        if (this.shouldCompleteDrain(shard)) {
+          this.removeShard(shard.id);
+        }
       }
 
       return;
@@ -368,11 +421,24 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
       }
     }
 
-    // Get active shards
-    const activeShards = roomShards
-      .filter(shard => shard && shard.status() === 'active');
+    // Get active shards with available capacity
+    let activeShards = this.getAvailableShards(roomShards);
 
     if (activeShards.length === 0) {
+      if (autoCreate && this.canCreateShard(room, roomShards.length)) {
+        const newShard = await this.createShard(roomId);
+        if (newShard) {
+          return {
+            shardId: newShard.id,
+            url: newShard.url()
+          };
+        }
+      }
+
+      if (roomShards.some(shard => shard.status() === 'active')) {
+        return { error: `No shard capacity available for room ${roomId}` };
+      }
+
       return { error: `No active shards available for room ${roomId}` };
     }
 
@@ -412,6 +478,18 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
     };
   }
 
+  private getAvailableShards(shards: ShardInfo[]) {
+    return shards.filter(shard =>
+      shard
+      && shard.status() === 'active'
+      && shard.currentConnections() < shard.maxConnections()
+    );
+  }
+
+  private canCreateShard(room: RoomConfig, currentShardCount: number) {
+    return room.maxShards() === undefined || currentShardCount < room.maxShards()!;
+  }
+
   // Private methods
   private async createShard(
     roomId: string,
@@ -425,7 +503,8 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
     }
 
     // Generate shard ID
-    const shardId = `${roomId}:${Date.now()}-${Math.floor(Math.random() * 10000)}`;
+    const worldId = this.getWorldId();
+    const shardId = `${roomId}:${worldId}:${Date.now()}-${Math.floor(Math.random() * 10000)}`;
 
     // Generate URL from template
     const template = urlTemplate || this.defaultShardUrlTemplate();
@@ -438,6 +517,7 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
     const newShard = new ShardInfo();
     newShard.id = shardId;
     newShard.roomId.set(roomId);
+    newShard.worldId.set(worldId);
     newShard.url.set(url);
     newShard.maxConnections.set(max);
     newShard.currentConnections.set(0);
@@ -448,4 +528,24 @@ export class WorldRoom implements RoomInterceptorPacket, RoomOnJoin {
     this.shards()[shardId] = newShard;
     return newShard;
   }
+
+  private async ensureMinShards(roomId: string) {
+    const room = this.rooms()[roomId];
+    if (!room) {
+      return;
+    }
+
+    const currentShardCount = Object.values(this.shards())
+      .filter(shard => shard.roomId() === roomId)
+      .length;
+    const targetShardCount = room.minShards();
+
+    if (currentShardCount >= targetShardCount) {
+      return;
+    }
+
+    for (let i = currentShardCount; i < targetShardCount; i++) {
+      await this.createShard(roomId);
+    }
+  }
 }

package/examples/game/.vscode/launch.json

@@ -1,11 +0,0 @@
-{
-  "configurations": [
-    {
-      "type": "node",
-      "request": "attach",
-      "name": "PartyKit debugger",
-      "address": "localhost",
-      "port": 9229
-    }
-  ]
-}