refactor server to support room creation and state updates; add WebSocket server initialization

Co-authored-by: Copilot <copilot@github.com>

src/server.test.ts

@@ -1,6 +1,7 @@
-import { describe, it } from "node:test";
+import { describe, it, before, after } from "node:test";
 import assert from "node:assert";
-import { validateNickname, createRoomId } from "./server";
+import { validateNickname, createRoomId, createServer } from "./server.js";
+import { WebSocket, WebSocketServer } from "ws";
 
 describe("DiscoveryServer pure functions", () => {
   it("validateNickname accepts 1-16 chars", () => {
@@ -22,3 +23,137 @@ describe("DiscoveryServer pure functions", () => {
     assert.match(a, /^r\d+-\d+$/);
   });
 });
+
+describe("DiscoveryServer Room Logic", () => {
+  let wss: WebSocketServer;
+  const port = 3002;
+  const url = `ws://localhost:${port}`;
+
+  before(() => {
+    wss = createServer(port);
+  });
+
+  after(() => {
+    wss.close();
+  });
+
+  function createClient(): Promise<
+    WebSocket & { nextMessage(): Promise<any> }
+  > {
+    return new Promise((resolve) => {
+      const ws = new WebSocket(url);
+      const messages: any[] = [];
+      let messageResolver: ((msg: any) => void) | null = null;
+
+      ws.on("message", (data) => {
+        const msg = JSON.parse(data.toString());
+        if (messageResolver) {
+          messageResolver(msg);
+          messageResolver = null;
+        } else {
+          messages.push(msg);
+        }
+      });
+
+      const nextMessage = () => {
+        return new Promise<any>((res) => {
+          if (messages.length > 0) {
+            res(messages.shift());
+          } else {
+            messageResolver = res;
+          }
+        });
+      };
+
+      ws.on("open", () => {
+        Object.assign(ws, { nextMessage });
+        resolve(ws as any);
+      });
+    });
+  }
+
+  it("creates a room and assigns host", async () => {
+    const ws1 = await createClient();
+    await ws1.nextMessage(); // welcome
+
+    ws1.send(JSON.stringify({ type: "create_room" }));
+    const msg = await ws1.nextMessage();
+
+    assert.strictEqual(msg.type, "room_created");
+    assert.ok(msg.roomId);
+    ws1.close();
+  });
+
+  it("rejects 3rd player from joining a room", async () => {
+    const ws1 = await createClient();
+    const welcome1 = await ws1.nextMessage();
+
+    const ws2 = await createClient();
+    await ws2.nextMessage();
+
+    const ws3 = await createClient();
+    await ws3.nextMessage();
+
+    // ws1 creates room
+    ws1.send(JSON.stringify({ type: "create_room" }));
+    const createMsg = await ws1.nextMessage();
+    const roomId = createMsg.roomId;
+    const hostId = welcome1.id;
+
+    // ws2 joins room
+    ws2.send(JSON.stringify({ type: "join_room", targetId: hostId }));
+    const joinMsg = await ws2.nextMessage();
+    assert.strictEqual(joinMsg.type, "room_joined");
+
+    // ignore peer_joined on ws1
+    await ws1.nextMessage();
+
+    // ws3 tries to join
+    ws3.send(JSON.stringify({ type: "join_room", targetId: hostId }));
+    const rejectMsg = await ws3.nextMessage();
+    assert.strictEqual(rejectMsg.type, "error");
+    assert.strictEqual(rejectMsg.message, "room_full");
+
+    ws1.close();
+    ws2.close();
+    ws3.close();
+  });
+
+  it("cleans up room on disconnect", async () => {
+    const ws1 = await createClient();
+    const welcome1 = await ws1.nextMessage();
+
+    const ws2 = await createClient();
+    await ws2.nextMessage();
+
+    // ws1 creates room
+    ws1.send(JSON.stringify({ type: "create_room" }));
+    await ws1.nextMessage(); // room_created
+
+    // ws2 joins
+    ws2.send(JSON.stringify({ type: "join_room", targetId: welcome1.id }));
+    await ws2.nextMessage(); // room_joined
+    await ws1.nextMessage(); // peer_joined
+
+    // ws1 disconnects
+    ws1.close();
+
+    // ws2 receives peer_left
+    const leftMsg = await ws2.nextMessage();
+    assert.strictEqual(leftMsg.type, "peer_left");
+    ws2.close();
+  });
+
+  it("validates room join target", async () => {
+    const ws1 = await createClient();
+    await ws1.nextMessage();
+
+    ws1.send(JSON.stringify({ type: "join_room", targetId: "nonexistent" }));
+    const rejectMsg = await ws1.nextMessage();
+
+    assert.strictEqual(rejectMsg.type, "error");
+    assert.strictEqual(rejectMsg.message, "invalid_target");
+
+    ws1.close();
+  });
+});

src/server.ts

@@ -33,12 +33,13 @@ const isMain =
   process.argv[1]?.endsWith("server.ts") ||
   process.argv[1]?.endsWith("server.js");
 
-if (isMain) {
+export function createServer(port: number): WebSocketServer {
-  const wss = new WebSocketServer({ port: 3001 });
+  const wss = new WebSocketServer({ port });
   const players = new Map<string, Player>();
   const rooms = new Map<string, Room>();
 
   let nextPlayerId = 1;
+  let serverTick = 0;
 
   wss.on("connection", (ws) => {
     const playerId = `p${nextPlayerId++}`;
@@ -71,7 +72,6 @@ if (isMain) {
             break;
 
           case "create_room": {
-            // Player creates a room and becomes host
             if (player.roomId) {
               ws.send(
                 JSON.stringify({ type: "error", message: "already_in_room" }),
@@ -87,7 +87,6 @@ if (isMain) {
           }
 
           case "join_room": {
-            // Player joins another player's room
             const targetId = msg.targetId as string | undefined;
             if (!targetId || !players.has(targetId)) {
               ws.send(
@@ -116,11 +115,9 @@ if (isMain) {
               ws.send(JSON.stringify({ type: "error", message: "room_full" }));
               break;
             }
-            // Add player to room
             room.players.push(playerId);
             player.roomId = target.roomId;
 
-            // Notify both players
             ws.send(
               JSON.stringify({
                 type: "room_joined",
@@ -140,6 +137,24 @@ if (isMain) {
             break;
           }
 
+          case "state_update": {
+            if (!player.roomId) break;
+            const room = rooms.get(player.roomId);
+            if (!room) break;
+
+            serverTick++;
+            msg.tick = serverTick;
+            msg.playerId = playerId;
+
+            const outMsg = JSON.stringify(msg);
+            for (const pId of room.players) {
+              if (pId !== playerId) {
+                players.get(pId)?.ws.send(outMsg);
+              }
+            }
+            break;
+          }
+
           default:
             ws.send(JSON.stringify({ type: "error", message: "unknown_type" }));
         }
@@ -149,12 +164,10 @@ if (isMain) {
     });
 
     ws.on("close", () => {
-      // Clean up room if player was in one
       if (player.roomId) {
         const room = rooms.get(player.roomId);
         if (room) {
           room.players = room.players.filter((p) => p !== playerId);
-          // Notify remaining players
           for (const pId of room.players) {
             const p = players.get(pId);
             if (p) {
@@ -167,7 +180,6 @@ if (isMain) {
               );
             }
           }
-          // Delete empty rooms
           if (room.players.length === 0) {
             rooms.delete(player.roomId);
           }
@@ -179,5 +191,10 @@ if (isMain) {
     ws.send(JSON.stringify({ type: "welcome", id: playerId }));
   });
 
+  return wss;
+}
+
+if (isMain) {
+  createServer(3001);
   console.log("DiscoveryServer listening on ws://localhost:3001");
 }