@signe/room 1.4.1 → 2.0.0

package/src/server.ts

@@ -61,19 +61,19 @@ export class Server implements Party.Server {
    * const server = new MyServer(new ServerIo("game"));
    * ```
    */
-  constructor(readonly room: Party.Room) {}
+  constructor(readonly room: Party.Room) { }
 
- /**
-   * @readonly
-   * @property {boolean} isHibernate - Indicates whether the server is in hibernate mode.
-   * 
-   * @example
-   * ```typescript
-   * if (!server.isHibernate) {
-   *   console.log("Server is active");
-   * }
-   * ```
-   */
+  /**
+    * @readonly
+    * @property {boolean} isHibernate - Indicates whether the server is in hibernate mode.
+    * 
+    * @example
+    * ```typescript
+    * if (!server.isHibernate) {
+    *   console.log("Server is active");
+    * }
+    * ```
+    */
   get isHibernate(): boolean {
     return !!this["options"]?.hibernate;
   }
@@ -82,6 +82,24 @@ export class Server implements Party.Server {
     return this.room.storage
   }
 
+  async send(conn: Party.Connection, obj: any, subRoom: any) {
+    obj = structuredClone(obj);
+    if (subRoom.interceptorPacket) {
+      const signal = this.getUsersProperty(subRoom);
+      const { publicId } = conn.state as any;
+      const user = signal?.()[publicId];
+      obj = await awaitReturn(subRoom["interceptorPacket"]?.(user, obj, conn));
+      if (obj === null) return;
+    }
+    conn.send(JSON.stringify(obj));
+  }
+
+  broadcast(obj: any, subRoom: any) {
+    for (let conn of this.room.getConnections()) {
+      this.send(conn, obj, subRoom);
+    }
+  }
+
   /**
    * @method onStart
    * @async
@@ -122,7 +140,7 @@ export class Server implements Party.Server {
       // Store valid publicIds from sessions
       const validPublicIds = new Set<string>();
       const expiredPublicIds = new Set<string>();
-      const SESSION_EXPIRY_TIME = options.sessionExpiryTime 
+      const SESSION_EXPIRY_TIME = options.sessionExpiryTime
       const now = Date.now();
 
       for (const [key, session] of sessions) {
@@ -130,15 +148,15 @@ export class Server implements Party.Server {
         if (!key.startsWith('session:')) continue;
 
         const privateId = key.replace('session:', '');
-        const typedSession = session as {publicId: string, created: number, connected: boolean};
-        
+        const typedSession = session as { publicId: string, created: number, connected: boolean };
+
         // Check if session should be deleted based on:
         // 1. Connection is not active
         // 2. Session is marked as disconnected
         // 3. Session is older than expiry time
-        if (!activePrivateIds.has(privateId) && 
-            !typedSession.connected && 
-            (now - typedSession.created) > SESSION_EXPIRY_TIME) {
+        if (!activePrivateIds.has(privateId) &&
+          !typedSession.connected &&
+          (now - typedSession.created) > SESSION_EXPIRY_TIME) {
           // Delete expired session
           await this.deleteSession(privateId);
           expiredPublicIds.add(typedSession.publicId);
@@ -159,7 +177,7 @@ export class Server implements Party.Server {
           }
         }
       }
-     
+
     } catch (error) {
       console.error('Error in garbage collector:', error);
     }
@@ -229,11 +247,12 @@ export class Server implements Party.Server {
         return;
       }
       const packet = buildObject(values, instance.$memoryAll);
-      this.room.broadcast(
-        JSON.stringify({
+      this.broadcast(
+        {
           type: "sync",
           value: packet,
-        })
+        },
+        instance
       );
       values.clear();
     }
@@ -247,7 +266,7 @@ export class Server implements Party.Server {
       for (let [path, value] of values) {
         const _instance =
           path == "." ? instance : getByPath(instance, path);
-        const itemValue = createStatesSnapshot(_instance); 
+        const itemValue = createStatesSnapshot(_instance);
         if (value == DELETE_TOKEN) {
           await this.room.storage.delete(path);
         } else {
@@ -327,17 +346,17 @@ export class Server implements Party.Server {
     return meta?.get("users")
   }
 
-  private async getSession(privateId: string): Promise<{publicId: string, state?: any, created?: number, connected?: boolean} | null> {
+  private async getSession(privateId: string): Promise<{ publicId: string, state?: any, created?: number, connected?: boolean } | 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;
+      return session as { publicId: string, state?: any, created: number, connected: boolean } | null;
     } catch (e) {
       return null;
     }
   }
 
-  private async saveSession(privateId: string, data: {publicId: string, state?: any, created?: number, connected?: boolean}) {
+  private async saveSession(privateId: string, data: { publicId: string, state?: any, created?: number, connected?: boolean }) {
     const sessionData = {
       ...data,
       created: data.created || Date.now(),
@@ -357,23 +376,7 @@ export class Server implements Party.Server {
     await this.room.storage.delete(`session:${privateId}`);
   }
 
-  /**
-   * @method onConnect
-   * @async
-   * @param {Party.Connection} conn - The connection object for the new user.
-   * @param {Party.ConnectionContext} ctx - The context of the connection.
-   * @description Handles a new user connection, creates a user object, and sends initial sync data.
-   * @returns {Promise<void>}
-   * 
-   * @example
-   * ```typescript
-   * server.onConnect = async (conn, ctx) => {
-   *   await server.onConnect(conn, ctx);
-   *   console.log("New user connected:", conn.id);
-   * };
-   * ```
-   */
-  async onConnect(conn: Party.Connection, ctx: Party.ConnectionContext) {
+  async onConnectClient(conn: Party.Connection, ctx: Party.ConnectionContext) {
     const subRoom = await this.getSubRoom({
       getMemoryAll: true,
     })
@@ -397,7 +400,7 @@ export class Server implements Party.Server {
     }
 
     // Check for existing session
-    const existingSession = await this.getSession(conn.id) 
+    const existingSession = await this.getSession(conn.id)
 
     // Generate IDs
     const publicId = existingSession?.publicId || generateShortUUID();
@@ -408,7 +411,7 @@ export class Server implements Party.Server {
 
     if (signal) {
       const { classType } = signal.options;
-     
+
       // Restore state if exists
       if (!existingSession?.publicId) {
         user = isClass(classType) ? new classType() : classType(conn, ctx);
@@ -416,7 +419,7 @@ export class Server implements Party.Server {
         const snapshot = createStatesSnapshot(user);
         this.room.storage.put(`${usersPropName}.${publicId}`, snapshot);
       }
-      
+
       // Only store new session if it doesn't exist
       if (!existingSession) {
         await this.saveSession(conn.id, {
@@ -430,40 +433,82 @@ export class Server implements Party.Server {
 
     // Call the room's onJoin method if it exists
     await awaitReturn(subRoom["onJoin"]?.(user, conn, ctx));
-    
+
     // Store both IDs in connection state
-    conn.setState({ publicId });
+    conn.setState({
+      ...conn.state,
+      publicId
+    });
 
     // Send initial sync data with both IDs to the new connection
-    conn.send(
-      JSON.stringify({
-        type: "sync",
-        value: {
-          pId: publicId,
-          ...subRoom.$memoryAll,
-        },
-      })
-    );
+    this.send(conn, {
+      type: "sync",
+      value: {
+        pId: publicId,
+        ...subRoom.$memoryAll,
+      },
+    }, subRoom)
   }
 
   /**
-   * @method onMessage
+   * @method onConnect
    * @async
-   * @param {string} message - The message received from a user.
-   * @param {Party.Connection} sender - The connection object of the sender.
-   * @description Processes incoming messages and triggers corresponding actions in the sub-room.
+   * @param {Party.Connection} conn - The connection object for the new user.
+   * @param {Party.ConnectionContext} ctx - The context of the connection.
+   * @description Handles a new user connection, creates a user object, and sends initial sync data.
    * @returns {Promise<void>}
    * 
    * @example
    * ```typescript
-   * server.onMessage = async (message, sender) => {
-   *   await server.onMessage(message, sender);
-   *   console.log("Message processed from:", sender.id);
+   * server.onConnect = async (conn, ctx) => {
+   *   await server.onConnect(conn, ctx);
+   *   console.log("New user connected:", conn.id);
    * };
    * ```
    */
+  async onConnect(conn: Party.Connection, ctx: Party.ConnectionContext) {
+    if (ctx.request?.headers.has('x-shard-id')) {
+      this.onConnectShard(conn, ctx);
+    }
+    else {
+      await this.onConnectClient(conn, ctx);
+    }
+  }
 
+  /**
+   * @method onConnectShard
+   * @private
+   * @param {Party.Connection} conn - The connection object for the new shard.
+   * @param {Party.ConnectionContext} ctx - The context of the shard connection.
+   * @description Handles a new shard connection, setting up the necessary state.
+   * @returns {void}
+   */
+  onConnectShard(conn: Party.Connection, ctx: Party.ConnectionContext) {
+    // Set shard metadata in connection state
+    const shardId = ctx.request?.headers.get('x-shard-id') || 'unknown-shard';
+    conn.setState({
+      shard: true,
+      shardId,
+      clients: new Map() // Track clients connected through this shard
+    });
+  }
+
+  /**
+   * @method onMessage
+   * @async
+   * @param {string} message - The message received from a user or shard.
+   * @param {Party.Connection} sender - The connection object of the sender.
+   * @description Processes incoming messages, handling differently based on if sender is shard or client.
+   * @returns {Promise<void>}
+   */
   async onMessage(message: string, sender: Party.Connection) {
+    // Check if message is from a shard
+    if (sender.state && (sender.state as any).shard) {
+      await this.handleShardMessage(message, sender);
+      return;
+    }
+
+    // Regular client message handling
     let json
     try {
       json = JSON.parse(message)
@@ -471,16 +516,19 @@ export class Server implements Party.Server {
     catch (e) {
       return;
     }
-     // Validate incoming messages
+
+    // Validate incoming messages
     const result = Message.safeParse(json);
     if (!result.success) {
       return;
     }
+
     const subRoom = await this.getSubRoom()
+
     // Check room guards
     const roomGuards = subRoom.constructor['_roomGuards'] || [];
     for (const guard of roomGuards) {
-      const isAuthorized = await guard(sender, result.data.value);
+      const isAuthorized = await guard(sender, result.data.value, this.room);
       if (!isAuthorized) {
         return;
       }
@@ -520,6 +568,214 @@ export class Server implements Party.Server {
     }
   }
 
+  /**
+   * @method handleShardMessage
+   * @private
+   * @async
+   * @param {string} message - The message received from a shard.
+   * @param {Party.Connection} shardConnection - The connection object of the shard.
+   * @description Processes messages from shards, extracting client information.
+   * @returns {Promise<void>}
+   */
+  private async handleShardMessage(message: string, shardConnection: Party.Connection) {
+    let parsedMessage;
+    try {
+      parsedMessage = JSON.parse(message);
+    } catch (e) {
+      console.error("Error parsing shard message:", e);
+      return;
+    }
+
+    const shardState = shardConnection.state as any;
+    const clients = shardState.clients;
+
+    switch (parsedMessage.type) {
+      case 'shard.clientConnected':
+        // Handle new client connection through shard
+        await this.handleShardClientConnect(parsedMessage, shardConnection);
+        break;
+
+      case 'shard.clientMessage':
+        // Handle message from a client through shard
+        await this.handleShardClientMessage(parsedMessage, shardConnection);
+        break;
+
+      case 'shard.clientDisconnected':
+        // Handle client disconnection through shard
+        await this.handleShardClientDisconnect(parsedMessage, shardConnection);
+        break;
+
+      default:
+        console.warn(`Unknown shard message type: ${parsedMessage.type}`);
+    }
+  }
+
+  /**
+   * @method handleShardClientConnect
+   * @private
+   * @async
+   * @param {Object} message - The client connection message from a shard.
+   * @param {Party.Connection} shardConnection - The connection object of the shard.
+   * @description Handles a new client connection via a shard.
+   * @returns {Promise<void>}
+   */
+  private async handleShardClientConnect(message: any, shardConnection: Party.Connection) {
+    const { privateId, connectionInfo } = message;
+    const shardState = shardConnection.state as any;
+
+    // Create a virtual connection context for the client
+    const virtualContext: Party.ConnectionContext = {
+      request: {
+        headers: new Headers({
+          'x-forwarded-for': connectionInfo.ip,
+          'user-agent': connectionInfo.userAgent,
+          // Add other headers as needed
+        }),
+        method: 'GET',
+        url: ''
+      } as unknown as Party.Request
+    };
+
+    // Create a virtual connection for the client
+    const virtualConnection: Partial<Party.Connection> = {
+      id: privateId,
+      send: (data: string) => {
+        // Forward to the actual client through the shard
+        shardConnection.send(JSON.stringify({
+          targetClientId: privateId,
+          data
+        }));
+      },
+      state: {},
+      setState: (state: unknown) => {
+        // Store client state in the shard's client map
+        const clients = shardState.clients;
+        const currentState = clients.get(privateId) || {};
+        const mergedState = Object.assign({}, currentState, state as object);
+        clients.set(privateId, mergedState);
+
+        // Update our virtual connection's state reference
+        virtualConnection.state = clients.get(privateId);
+        return virtualConnection.state as any;
+      },
+      close: () => {
+        // Send close command to the shard
+        shardConnection.send(JSON.stringify({
+          type: 'shard.closeClient',
+          privateId
+        }));
+
+        // Clean up virtual connection
+        if (shardState.clients) {
+          shardState.clients.delete(privateId);
+        }
+      }
+    };
+
+    // Initialize the client's state in the shard state
+    if (!shardState.clients.has(privateId)) {
+      shardState.clients.set(privateId, {});
+    }
+
+    // Now handle this virtual connection as a regular client connection
+    await this.onConnectClient(virtualConnection as Party.Connection, virtualContext);
+  }
+
+  /**
+   * @method handleShardClientMessage
+   * @private
+   * @async
+   * @param {Object} message - The client message from a shard.
+   * @param {Party.Connection} shardConnection - The connection object of the shard.
+   * @description Handles a message from a client via a shard.
+   * @returns {Promise<void>}
+   */
+  private async handleShardClientMessage(message: any, shardConnection: Party.Connection) {
+    const { privateId, publicId, payload } = message;
+    const shardState = shardConnection.state as any;
+    const clients = shardState.clients;
+
+    // Get or create virtual connection for this client
+    if (!clients.has(privateId)) {
+      console.warn(`Received message from unknown client ${privateId}, creating virtual connection`);
+      clients.set(privateId, { publicId });
+    }
+
+    // Create a virtual connection for the client
+    const virtualConnection: Partial<Party.Connection> = {
+      id: privateId,
+      send: (data: string) => {
+        // Forward to the actual client through the shard
+        shardConnection.send(JSON.stringify({
+          targetClientId: privateId,
+          data
+        }));
+      },
+      state: clients.get(privateId),
+      setState: (state: unknown) => {
+        const currentState = clients.get(privateId) || {};
+        const mergedState = Object.assign({}, currentState, state as object);
+        clients.set(privateId, mergedState);
+        virtualConnection.state = clients.get(privateId);
+        return virtualConnection.state as any;
+      },
+      close: () => {
+        shardConnection.send(JSON.stringify({
+          type: 'shard.closeClient',
+          privateId
+        }));
+
+        if (shardState.clients) {
+          shardState.clients.delete(privateId);
+        }
+      }
+    };
+
+    // Process the payload using the regular message handler
+    const payloadString = typeof payload === 'string' ? payload : JSON.stringify(payload);
+    await this.onMessage(payloadString, virtualConnection as Party.Connection);
+  }
+
+  /**
+   * @method handleShardClientDisconnect
+   * @private
+   * @async
+   * @param {Object} message - The client disconnection message from a shard.
+   * @param {Party.Connection} shardConnection - The connection object of the shard.
+   * @description Handles a client disconnection via a shard.
+   * @returns {Promise<void>}
+   */
+  private async handleShardClientDisconnect(message: any, shardConnection: Party.Connection) {
+    const { privateId, publicId } = message;
+    const shardState = shardConnection.state as any;
+    const clients = shardState.clients;
+
+    // Get client state
+    const clientState = clients.get(privateId);
+    if (!clientState) {
+      console.warn(`Disconnection for unknown client ${privateId}`);
+      return;
+    }
+
+    // Create a virtual connection for the client one last time
+    const virtualConnection: Partial<Party.Connection> = {
+      id: privateId,
+      send: () => { }, // No-op since client is disconnecting
+      state: clientState,
+      setState: () => {
+        // No-op since client is disconnecting
+        return {} as any;
+      },
+      close: () => { }
+    };
+
+    // Handle disconnection with the regular onClose handler
+    await this.onClose(virtualConnection as Party.Connection);
+
+    // Clean up
+    clients.delete(privateId);
+  }
+
   /**
    * @method onClose
    * @async
@@ -560,12 +816,10 @@ export class Server implements Party.Server {
     await this.updateSessionConnection(privateId, false);
 
     // Broadcast user disconnection
-    this.room.broadcast(
-      JSON.stringify({
-        type: "user_disconnected",
-        value: { publicId }
-      })
-    );
+    this.broadcast({
+      type: "user_disconnected",
+      value: { publicId }
+    }, subRoom);
   }
 
   async onAlarm() {
@@ -578,26 +832,288 @@ export class Server implements Party.Server {
     await awaitReturn(subRoom["onError"]?.(connection, error));
   }
 
+  /**
+   * @method onRequest
+   * @async
+   * @param {Party.Request} req - The HTTP request to handle
+   * @description Handles HTTP requests, either directly from clients or forwarded by shards
+   * @returns {Promise<Response>} The response to return to the client
+   */
   async onRequest(req: Party.Request) {
-    const subRoom = await this.getSubRoom()
+    // Check if the request is coming from a shard
+    const isFromShard = req.headers.has('x-forwarded-by-shard');
+    const shardId = req.headers.get('x-shard-id');
+
+    if (isFromShard) {
+      return this.handleShardRequest(req, shardId);
+    }
+
+    // Handle regular client request
+    return this.handleDirectRequest(req);
+  }
+
+  /**
+   * @method handleDirectRequest
+   * @private
+   * @async
+   * @param {Party.Request} req - The HTTP request received directly from a client
+   * @description Processes requests received directly from clients
+   * @returns {Promise<Response>} The response to return to the client
+   */
+  private async handleDirectRequest(req: Party.Request): Promise<Response> {
+    const subRoom = await this.getSubRoom();
     const res = (body: any, status: number) => {
       return new Response(JSON.stringify(body), { status });
-    }
+    };
+
     if (!subRoom) {
       return res({
         error: "Not found"
       }, 404);
     }
 
-    const response = await awaitReturn(subRoom["onRequest"]?.(req, this.room));
-    if (!response) {
+    // First try to match using the registered @Request handlers
+    const response = await this.tryMatchRequestHandler(req, subRoom);
+    if (response) {
+      return response;
+    }
+
+    // Fall back to the legacy onRequest method if no handler matched
+    const legacyResponse = await awaitReturn(subRoom["onRequest"]?.(req, this.room));
+    if (!legacyResponse) {
       return res({
         error: "Not found"
       }, 404);
     }
-    if (response instanceof Response) {
-      return response;
+    if (legacyResponse instanceof Response) {
+      return legacyResponse;
+    }
+    return res(legacyResponse, 200);
+  }
+
+  /**
+   * @method tryMatchRequestHandler
+   * @private
+   * @async
+   * @param {Party.Request} req - The HTTP request to handle
+   * @param {Object} subRoom - The room instance
+   * @description Attempts to match the request to a registered @Request handler
+   * @returns {Promise<Response | null>} The response or null if no handler matched
+   */
+  private async tryMatchRequestHandler(req: Party.Request, subRoom: any): Promise<Response | null> {
+    const requestHandlers = subRoom.constructor["_requestMetadata"];
+    if (!requestHandlers) {
+      return null;
+    }
+
+    const url = new URL(req.url);
+    const method = req.method;
+    let pathname = url.pathname;
+
+    pathname = '/' + pathname.split('/').slice(4).join('/');
+
+    // Check each registered handler
+    for (const [routeKey, handler] of requestHandlers.entries()) {
+      const firstColonIndex = routeKey.indexOf(':');
+      const handlerMethod = routeKey.substring(0, firstColonIndex);
+      const handlerPath = routeKey.substring(firstColonIndex + 1);
+
+      // Check if method matches
+      if (handlerMethod !== method) {
+        continue;
+      }
+
+      // Simple path matching (could be enhanced with path params)
+      if (this.pathMatches(pathname, handlerPath)) {
+        // Extract path params if any
+        const params = this.extractPathParams(pathname, handlerPath);
+        // Check request guards if they exist
+        const guards = subRoom.constructor['_actionGuards']?.get(handler.key) || [];
+        for (const guard of guards) {
+          const isAuthorized = await guard(null, req, this.room);
+          if (isAuthorized instanceof Response) {
+            return isAuthorized;
+          }
+          if (!isAuthorized) {
+            return new Response(JSON.stringify({ error: "Unauthorized" }), { status: 403 });
+          }
+        }
+
+        // Validate request body if needed
+        let bodyData = null;
+        if (handler.bodyValidation && ['POST', 'PUT', 'PATCH'].includes(method)) {
+          try {
+            const contentType = req.headers.get('content-type') || '';
+            if (contentType.includes('application/json')) {
+              const body = await req.json();
+              const validation = handler.bodyValidation.safeParse(body);
+              if (!validation.success) {
+                return new Response(
+                  JSON.stringify({ error: "Invalid request body", details: validation.error }),
+                  { status: 400 }
+                );
+              }
+              bodyData = validation.data;
+            }
+          } catch (error) {
+            return new Response(
+              JSON.stringify({ error: "Failed to parse request body" }),
+              { status: 400 }
+            );
+          }
+        }
+
+        // Execute the handler method
+        try {
+          const result = await awaitReturn(
+            subRoom[handler.key](req, bodyData, params, this.room)
+          );
+
+          if (result instanceof Response) {
+            return result;
+          }
+
+          return new Response(
+            typeof result === 'string' ? result : JSON.stringify(result),
+            {
+              status: 200,
+              headers: { 'Content-Type': typeof result === 'string' ? 'text/plain' : 'application/json' }
+            }
+          );
+        } catch (error) {
+          console.error('Error executing request handler:', error);
+          return new Response(
+            JSON.stringify({ error: "Internal server error" }),
+            { status: 500 }
+          );
+        }
+      }
     }
-    return res(response, 200);
+
+    return null;
+  }
+
+  /**
+   * @method pathMatches
+   * @private
+   * @param {string} requestPath - The path from the request
+   * @param {string} handlerPath - The path pattern from the handler
+   * @description Checks if a request path matches a handler path pattern
+   * @returns {boolean} True if the paths match
+   */
+  private pathMatches(requestPath: string, handlerPath: string): boolean {
+    // Convert handler path pattern to regex
+    // Replace :param with named capture groups
+    const pathRegexString = handlerPath
+      .replace(/\//g, '\\/') // Escape slashes
+      .replace(/:([^\/]+)/g, '([^/]+)'); // Convert :params to capture groups
+
+    const pathRegex = new RegExp(`^${pathRegexString}$`);
+    return pathRegex.test(requestPath);
+  }
+
+  /**
+   * @method extractPathParams
+   * @private
+   * @param {string} requestPath - The path from the request
+   * @param {string} handlerPath - The path pattern from the handler
+   * @description Extracts path parameters from the request path based on the handler pattern
+   * @returns {Object} An object containing the path parameters
+   */
+  private extractPathParams(requestPath: string, handlerPath: string): Record<string, string> {
+    const params: Record<string, string> = {};
+
+    // Extract parameter names from handler path
+    const paramNames: string[] = [];
+    handlerPath.split('/').forEach(segment => {
+      if (segment.startsWith(':')) {
+        paramNames.push(segment.substring(1));
+      }
+    });
+
+    // Extract parameter values from request path
+    const pathRegexString = handlerPath
+      .replace(/\//g, '\\/') // Escape slashes
+      .replace(/:([^\/]+)/g, '([^/]+)'); // Convert :params to capture groups
+
+    const pathRegex = new RegExp(`^${pathRegexString}$`);
+    const matches = requestPath.match(pathRegex);
+
+    if (matches && matches.length > 1) {
+      // Skip the first match (the full string)
+      for (let i = 0; i < paramNames.length; i++) {
+        params[paramNames[i]] = matches[i + 1];
+      }
+    }
+
+    return params;
+  }
+
+  /**
+   * @method handleShardRequest
+   * @private
+   * @async
+   * @param {Party.Request} req - The HTTP request forwarded by a shard
+   * @param {string | null} shardId - The ID of the shard that forwarded the request
+   * @description Processes requests forwarded by shards, preserving client context
+   * @returns {Promise<Response>} The response to return to the shard (which will forward it to the client)
+   */
+  private async handleShardRequest(req: Party.Request, shardId: string | null): Promise<Response> {
+    const subRoom = await this.getSubRoom();
+
+    if (!subRoom) {
+      return new Response(JSON.stringify({ error: "Not found" }), { status: 404 });
+    }
+
+    // Create a context that preserves original client information
+    const originalClientIp = req.headers.get('x-original-client-ip');
+    const enhancedReq = this.createEnhancedRequest(req, originalClientIp);
+
+    try {
+      // First try to match using the registered @Request handlers
+      const response = await this.tryMatchRequestHandler(enhancedReq, subRoom);
+      if (response) {
+        return response;
+      }
+
+      // Fall back to the legacy onRequest handler
+      const legacyResponse = await awaitReturn(subRoom["onRequest"]?.(enhancedReq, this.room));
+
+      if (!legacyResponse) {
+        return new Response(JSON.stringify({ error: "Not found" }), { status: 404 });
+      }
+
+      if (legacyResponse instanceof Response) {
+        return legacyResponse;
+      }
+
+      return new Response(JSON.stringify(legacyResponse), { status: 200 });
+    } catch (error) {
+      console.error(`Error processing request from shard ${shardId}:`, error);
+      return new Response(JSON.stringify({ error: "Internal server error" }), { status: 500 });
+    }
+  }
+
+  /**
+   * @method createEnhancedRequest
+   * @private
+   * @param {Party.Request} originalReq - The original request received from the shard
+   * @param {string | null} originalClientIp - The original client IP, if available
+   * @description Creates an enhanced request object that preserves the original client context
+   * @returns {Party.Request} The enhanced request object
+   */
+  private createEnhancedRequest(originalReq: Party.Request, originalClientIp: string | null): Party.Request {
+    // Clone the original request to avoid mutating it
+    const clonedReq = originalReq.clone();
+
+    // Add a custom property to the request to indicate it came via a shard
+    (clonedReq as any).viaShard = true;
+
+    // If we have the original client IP, we can use it for things like rate limiting or geolocation
+    if (originalClientIp) {
+      (clonedReq as any).originalClientIp = originalClientIp;
+    }
+
+    return clonedReq;
   }
 }