@quenty/studio-bridge 0.7.1 → 0.8.0

package/src/server/studio-bridge-server.test.ts

@@ -8,11 +8,10 @@
 
 import { describe, it, expect, vi, afterEach } from 'vitest';
 import { WebSocket } from 'ws';
-import { StudioBridgeServer, type StudioBridgePhase } from './studio-bridge-server.js';
-
-// ---------------------------------------------------------------------------
-// Mocks — replace external side-effects with no-ops
-// ---------------------------------------------------------------------------
+import {
+  StudioBridgeServer,
+  type StudioBridgePhase,
+} from './studio-bridge-server.js';
 
 vi.mock('@quenty/nevermore-template-helpers', () => ({
   BuildContext: {
@@ -42,13 +41,9 @@ vi.mock('../process/studio-process-manager.js', () => ({
   findPluginsFolder: vi.fn(() => '/fake/plugins'),
 }));
 
-// ---------------------------------------------------------------------------
-// Helpers
-// ---------------------------------------------------------------------------
-
 /**
- * Connect a WebSocket client to the server and perform the hello/welcome
- * handshake, returning the connected client.
+ * Connect a WebSocket client to the server and send a register message,
+ * returning the connected client.
  */
 async function connectAndHandshake(
   port: number,
@@ -61,20 +56,22 @@ async function connectAndHandshake(
     ws.on('error', reject);
   });
 
-  // Send hello
-  ws.send(JSON.stringify({ type: 'hello', sessionId, payload: { sessionId } }));
-
-  // Wait for welcome
-  await new Promise<void>((resolve) => {
-    ws.on('message', (raw) => {
-      const data = JSON.parse(
-        typeof raw === 'string' ? raw : raw.toString('utf-8')
-      );
-      if (data.type === 'welcome') {
-        resolve();
-      }
-    });
-  });
+  ws.send(
+    JSON.stringify({
+      type: 'register',
+      sessionId,
+      payload: {
+        pluginVersion: '1.0.0',
+        instanceId: 'inst-1',
+        placeName: 'TestPlace',
+        state: 'Edit',
+        capabilities: ['execute', 'queryState'],
+      },
+    })
+  );
+
+  // Allow the server to process the register message
+  await new Promise((r) => setTimeout(r, 20));
 
   return ws;
 }
@@ -85,7 +82,11 @@ async function connectAndHandshake(
  * and the port.
  */
 async function createReadyServer(
-  options: { sessionId?: string; timeoutMs?: number; onPhase?: (p: StudioBridgePhase) => void } = {}
+  options: {
+    sessionId?: string;
+    timeoutMs?: number;
+    onPhase?: (p: StudioBridgePhase) => void;
+  } = {}
 ) {
   const sessionId = options.sessionId ?? 'test-session';
   const server = new StudioBridgeServer({
@@ -98,16 +99,14 @@ async function createReadyServer(
   // Start in background — it will wait for handshake
   const startPromise = server.startAsync();
 
-  // We need to discover the port the server is listening on. The WSS is
-  // created inside startAsync and we can't access it directly. However, we
+  // We need to discover the port the server is listening on. The HTTP server
+  // is created inside startAsync and we can't access it directly. However, we
   // know the mock for launchStudioAsync will resolve immediately, so the
   // server will be waiting for a handshake. We just need the port.
   // Access it via the private field (acceptable in tests).
-  // Wait a tick for the WSS to be created
+  // Wait a tick for the HTTP server to be created
   await new Promise((r) => setTimeout(r, 50));
-  const wss = (server as any)._wss;
-  const addr = wss.address();
-  const port: number = addr.port;
+  const port: number = (server as any)._port;
 
   // Simulate plugin connecting and performing handshake
   const client = await connectAndHandshake(port, sessionId);
@@ -118,10 +117,6 @@ async function createReadyServer(
   return { server, client, port, sessionId };
 }
 
-// ---------------------------------------------------------------------------
-// Tests
-// ---------------------------------------------------------------------------
-
 describe('StudioBridgeServer', () => {
   let server: StudioBridgeServer | undefined;
   let client: WebSocket | undefined;
@@ -138,10 +133,6 @@ describe('StudioBridgeServer', () => {
     server = undefined;
   });
 
-  // -----------------------------------------------------------------------
-  // State machine guards
-  // -----------------------------------------------------------------------
-
   describe('state guards', () => {
     it('throws when executeAsync is called before startAsync', async () => {
       server = new StudioBridgeServer({ placePath: '/fake/place.rbxl' });
@@ -171,64 +162,17 @@ describe('StudioBridgeServer', () => {
     });
   });
 
-  // -----------------------------------------------------------------------
-  // Handshake
-  // -----------------------------------------------------------------------
-
   describe('handshake', () => {
-    it('accepts hello with correct session ID and sends welcome', async () => {
+    it('accepts register with correct session ID', async () => {
       const ready = await createReadyServer({ sessionId: 'my-session' });
       server = ready.server;
       client = ready.client;
 
-      // If we got here, the handshake succeeded (connectAndHandshake waits
-      // for the welcome message)
+      // If we got here the server accepted the register handshake and
+      // resolved startAsync — connectAndHandshake yields the connected client.
       expect(true).toBe(true);
     });
 
-    it('rejects hello with wrong session ID and closes connection', async () => {
-      const sessionId = 'correct-session';
-      server = new StudioBridgeServer({
-        placePath: '/fake/place.rbxl',
-        sessionId,
-        timeoutMs: 1_000,
-      });
-
-      const startPromise = server.startAsync();
-
-      await new Promise((r) => setTimeout(r, 50));
-      const wss = (server as any)._wss;
-      const port: number = wss.address().port;
-
-      // Connect with correct path but wrong session ID in message
-      const ws = new WebSocket(`ws://localhost:${port}/${sessionId}`);
-      await new Promise<void>((resolve, reject) => {
-        ws.on('open', resolve);
-        ws.on('error', reject);
-      });
-
-      // Track if server closes the connection
-      const closedPromise = new Promise<void>((resolve) => {
-        ws.on('close', () => resolve());
-      });
-
-      ws.send(
-        JSON.stringify({
-          type: 'hello',
-          sessionId: 'wrong-session',
-          payload: { sessionId: 'wrong-session' },
-        })
-      );
-
-      // Server should close the connection
-      await closedPromise;
-
-      // Server should NOT accept — startAsync should time out
-      await expect(startPromise).rejects.toThrow('Timed out waiting for Studio plugin handshake');
-
-      // Server already stopped on failure (startAsync catch cleans up)
-    });
-
     it('rejects connection to wrong WebSocket path', async () => {
       const sessionId = 'correct-session';
       server = new StudioBridgeServer({
@@ -240,8 +184,7 @@ describe('StudioBridgeServer', () => {
       const startPromise = server.startAsync();
 
       await new Promise((r) => setTimeout(r, 50));
-      const wss = (server as any)._wss;
-      const port: number = wss.address().port;
+      const port: number = (server as any)._port;
 
       // Connect with wrong path
       const ws = new WebSocket(`ws://localhost:${port}/wrong-path`);
@@ -258,14 +201,12 @@ describe('StudioBridgeServer', () => {
       expect(['error', 'rejected']).toContain(errorOrClose);
 
       // startAsync should time out
-      await expect(startPromise).rejects.toThrow('Timed out waiting for Studio plugin handshake');
+      await expect(startPromise).rejects.toThrow(
+        'Timed out waiting for Studio plugin handshake'
+      );
     });
   });
 
-  // -----------------------------------------------------------------------
-  // Script execution
-  // -----------------------------------------------------------------------
-
   describe('executeAsync', () => {
     it('sends execute message with sessionId and returns result on scriptComplete', async () => {
       const ready = await createReadyServer();
@@ -273,16 +214,21 @@ describe('StudioBridgeServer', () => {
       client = ready.client;
 
       // Listen for execute message from server
-      const executePromise = new Promise<{ script: string; sessionId: string }>((resolve) => {
-        client!.on('message', (raw) => {
-          const data = JSON.parse(
-            typeof raw === 'string' ? raw : raw.toString('utf-8')
-          );
-          if (data.type === 'execute') {
-            resolve({ script: data.payload.script, sessionId: data.sessionId });
-          }
-        });
-      });
+      const executePromise = new Promise<{ script: string; sessionId: string }>(
+        (resolve) => {
+          client!.on('message', (raw) => {
+            const data = JSON.parse(
+              typeof raw === 'string' ? raw : raw.toString('utf-8')
+            );
+            if (data.type === 'execute') {
+              resolve({
+                script: data.payload.script,
+                sessionId: data.sessionId,
+              });
+            }
+          });
+        }
+      );
 
       // Start execution
       const resultPromise = server.executeAsync({
@@ -308,117 +254,6 @@ describe('StudioBridgeServer', () => {
       expect(result.logs).toBe('');
     });
 
-    it('collects output messages into logs', async () => {
-      const ready = await createReadyServer();
-      server = ready.server;
-      client = ready.client;
-
-      const resultPromise = server.executeAsync({
-        scriptContent: 'print("test")',
-      });
-
-      // Wait for execute message to arrive
-      await new Promise<void>((resolve) => {
-        client!.on('message', (raw) => {
-          const data = JSON.parse(
-            typeof raw === 'string' ? raw : raw.toString('utf-8')
-          );
-          if (data.type === 'execute') resolve();
-        });
-      });
-
-      // Send output messages
-      client!.send(
-        JSON.stringify({
-          type: 'output',
-          sessionId: ready.sessionId,
-          payload: {
-            messages: [
-              { level: 'Print', body: 'Line 1' },
-              { level: 'Warning', body: 'Line 2' },
-            ],
-          },
-        })
-      );
-
-      client!.send(
-        JSON.stringify({
-          type: 'output',
-          sessionId: ready.sessionId,
-          payload: {
-            messages: [{ level: 'Error', body: 'Line 3' }],
-          },
-        })
-      );
-
-      // Complete
-      client!.send(
-        JSON.stringify({
-          type: 'scriptComplete',
-          sessionId: ready.sessionId,
-          payload: { success: true },
-        })
-      );
-
-      const result = await resultPromise;
-      expect(result.success).toBe(true);
-      expect(result.logs).toContain('Line 1');
-      expect(result.logs).toContain('Line 2');
-      expect(result.logs).toContain('Line 3');
-    });
-
-    it('calls onOutput for each output entry', async () => {
-      const ready = await createReadyServer();
-      server = ready.server;
-      client = ready.client;
-
-      const outputEntries: Array<{ level: string; body: string }> = [];
-
-      const resultPromise = server.executeAsync({
-        scriptContent: 'print("test")',
-        onOutput: (level, body) => {
-          outputEntries.push({ level, body });
-        },
-      });
-
-      await new Promise<void>((resolve) => {
-        client!.on('message', (raw) => {
-          const data = JSON.parse(
-            typeof raw === 'string' ? raw : raw.toString('utf-8')
-          );
-          if (data.type === 'execute') resolve();
-        });
-      });
-
-      client!.send(
-        JSON.stringify({
-          type: 'output',
-          sessionId: ready.sessionId,
-          payload: {
-            messages: [
-              { level: 'Print', body: 'hello' },
-              { level: 'Error', body: 'oops' },
-            ],
-          },
-        })
-      );
-
-      client!.send(
-        JSON.stringify({
-          type: 'scriptComplete',
-          sessionId: ready.sessionId,
-          payload: { success: true },
-        })
-      );
-
-      await resultPromise;
-
-      expect(outputEntries).toEqual([
-        { level: 'Print', body: 'hello' },
-        { level: 'Error', body: 'oops' },
-      ]);
-    });
-
     it('returns failure with error message on script error', async () => {
       const ready = await createReadyServer();
       server = ready.server;
@@ -473,7 +308,9 @@ describe('StudioBridgeServer', () => {
 
       const result = await resultPromise;
       expect(result.success).toBe(false);
-      expect(result.logs).toContain('Plugin disconnected before script completed');
+      expect(result.logs).toContain(
+        'Plugin disconnected before script completed'
+      );
     });
 
     it('times out when script takes too long', async () => {
@@ -528,10 +365,6 @@ describe('StudioBridgeServer', () => {
     });
   });
 
-  // -----------------------------------------------------------------------
-  // Multi-execution (reuse)
-  // -----------------------------------------------------------------------
-
   describe('multi-execution', () => {
     it('supports executing multiple scripts on the same session', async () => {
       const ready = await createReadyServer();
@@ -577,63 +410,8 @@ describe('StudioBridgeServer', () => {
       const r3 = await runOnce('print("third")');
       expect(r3.success).toBe(true);
     });
-
-    it('isolates logs between executions', async () => {
-      const ready = await createReadyServer();
-      server = ready.server;
-      client = ready.client;
-
-      const runWithOutput = async (outputBody: string) => {
-        const resultPromise = server!.executeAsync({
-          scriptContent: 'print("x")',
-        });
-
-        await new Promise<void>((resolve) => {
-          const handler = (raw: any) => {
-            const data = JSON.parse(
-              typeof raw === 'string' ? raw : raw.toString('utf-8')
-            );
-            if (data.type === 'execute') {
-              client!.off('message', handler);
-              resolve();
-            }
-          };
-          client!.on('message', handler);
-        });
-
-        client!.send(
-          JSON.stringify({
-            type: 'output',
-            sessionId: ready.sessionId,
-            payload: { messages: [{ level: 'Print', body: outputBody }] },
-          })
-        );
-
-        client!.send(
-          JSON.stringify({
-            type: 'scriptComplete',
-            sessionId: ready.sessionId,
-            payload: { success: true },
-          })
-        );
-
-        return resultPromise;
-      };
-
-      const r1 = await runWithOutput('output-from-first');
-      const r2 = await runWithOutput('output-from-second');
-
-      expect(r1.logs).toBe('output-from-first');
-      expect(r1.logs).not.toContain('output-from-second');
-      expect(r2.logs).toBe('output-from-second');
-      expect(r2.logs).not.toContain('output-from-first');
-    });
   });
 
-  // -----------------------------------------------------------------------
-  // Phase callbacks
-  // -----------------------------------------------------------------------
-
   describe('onPhase', () => {
     it('fires phase callbacks in order during lifecycle', async () => {
       const phases: StudioBridgePhase[] = [];
@@ -684,10 +462,6 @@ describe('StudioBridgeServer', () => {
     });
   });
 
-  // -----------------------------------------------------------------------
-  // Cleanup / stopAsync
-  // -----------------------------------------------------------------------
-
   describe('stopAsync', () => {
     it('sends shutdown message with sessionId to connected client', async () => {
       const ready = await createReadyServer();
@@ -702,7 +476,11 @@ describe('StudioBridgeServer', () => {
       await server.stopAsync();
 
       expect(sendSpy).toHaveBeenCalledWith(
-        JSON.stringify({ type: 'shutdown', sessionId: ready.sessionId, payload: {} })
+        JSON.stringify({
+          type: 'shutdown',
+          sessionId: ready.sessionId,
+          payload: {},
+        })
       );
     });
 
@@ -712,4 +490,264 @@ describe('StudioBridgeServer', () => {
       await server.stopAsync();
     });
   });
+
+  describe('heartbeat', () => {
+    it('silently accepts heartbeat messages after handshake', async () => {
+      const ready = await createReadyServer();
+      server = ready.server;
+      client = ready.client;
+
+      // Send a heartbeat message
+      client.send(
+        JSON.stringify({
+          type: 'heartbeat',
+          sessionId: ready.sessionId,
+          payload: {
+            uptimeMs: 5000,
+            state: 'Edit',
+            pendingRequests: 0,
+          },
+        })
+      );
+
+      // Give it time to process
+      await new Promise((r) => setTimeout(r, 50));
+
+      // Verify the server recorded the heartbeat
+      const lastHb = (server as any)._lastHeartbeatTimestamp;
+      expect(lastHb).toBeGreaterThan(0);
+    });
+
+    it('does not interfere with script execution', async () => {
+      const ready = await createReadyServer();
+      server = ready.server;
+      client = ready.client;
+
+      const resultPromise = server.executeAsync({
+        scriptContent: 'print("test")',
+      });
+
+      await new Promise<void>((resolve) => {
+        client!.on('message', (raw) => {
+          const data = JSON.parse(
+            typeof raw === 'string' ? raw : raw.toString('utf-8')
+          );
+          if (data.type === 'execute') resolve();
+        });
+      });
+
+      // Send a heartbeat during execution
+      client!.send(
+        JSON.stringify({
+          type: 'heartbeat',
+          sessionId: ready.sessionId,
+          payload: {
+            uptimeMs: 10000,
+            state: 'Edit',
+            pendingRequests: 1,
+          },
+        })
+      );
+
+      // Then complete the script
+      client!.send(
+        JSON.stringify({
+          type: 'scriptComplete',
+          sessionId: ready.sessionId,
+          payload: { success: true },
+        })
+      );
+
+      const result = await resultPromise;
+      expect(result.success).toBe(true);
+    });
+  });
+
+  describe('performActionAsync', () => {
+    /**
+     * Connect with v2 register, start the server, and return a ready state.
+     */
+    async function createV2ReadyServer(options?: {
+      sessionId?: string;
+      capabilities?: string[];
+    }) {
+      const sessionId = options?.sessionId ?? 'v2-action-session';
+      const capabilities = options?.capabilities ?? [
+        'execute',
+        'queryState',
+        'captureScreenshot',
+      ];
+
+      const srv = new StudioBridgeServer({
+        placePath: '/fake/place.rbxl',
+        sessionId,
+        timeoutMs: 5_000,
+      });
+
+      const startPromise = srv.startAsync();
+      await new Promise((r) => setTimeout(r, 50));
+      const port: number = (srv as any)._port;
+
+      const ws = new WebSocket(`ws://localhost:${port}/${sessionId}`);
+      await new Promise<void>((resolve, reject) => {
+        ws.on('open', resolve);
+        ws.on('error', reject);
+      });
+
+      ws.send(
+        JSON.stringify({
+          type: 'register',
+          sessionId,
+          payload: {
+            pluginVersion: '2.0.0',
+            instanceId: 'inst-action',
+            placeName: 'ActionTestPlace',
+            state: 'Edit',
+            capabilities,
+          },
+        })
+      );
+
+      await startPromise;
+
+      return { server: srv, client: ws, port, sessionId };
+    }
+
+    it('sends queryState action and resolves with stateResult', async () => {
+      const ready = await createV2ReadyServer();
+      server = ready.server;
+      client = ready.client;
+
+      // Listen for the queryState action from the server
+      const actionPromise = new Promise<Record<string, unknown>>((resolve) => {
+        client!.on('message', (raw) => {
+          const data = JSON.parse(
+            typeof raw === 'string' ? raw : raw.toString('utf-8')
+          );
+          if (data.type === 'queryState') {
+            resolve(data);
+          }
+        });
+      });
+
+      // Perform the action
+      const resultPromise = server.performActionAsync({
+        type: 'queryState',
+        payload: {},
+      });
+
+      // Wait for the queryState message to arrive at the mock plugin
+      const queryMsg = await actionPromise;
+      expect(queryMsg.type).toBe('queryState');
+      expect(queryMsg.sessionId).toBe(ready.sessionId);
+      expect(typeof queryMsg.requestId).toBe('string');
+
+      // Respond from the mock plugin
+      client!.send(
+        JSON.stringify({
+          type: 'stateResult',
+          sessionId: ready.sessionId,
+          requestId: queryMsg.requestId,
+          payload: {
+            state: 'Edit',
+            placeId: 12345,
+            placeName: 'ActionTestPlace',
+            gameId: 67890,
+          },
+        })
+      );
+
+      const result = await resultPromise;
+      expect(result.type).toBe('stateResult');
+      expect((result as any).payload.state).toBe('Edit');
+      expect((result as any).payload.placeId).toBe(12345);
+    });
+
+    it('rejects when plugin does not support requested capability', async () => {
+      // Connect with only 'execute' capability — no 'queryState'
+      const ready = await createV2ReadyServer({ capabilities: ['execute'] });
+      server = ready.server;
+      client = ready.client;
+
+      await expect(
+        server.performActionAsync({ type: 'queryState', payload: {} })
+      ).rejects.toThrow('Plugin does not support capability: queryState');
+    });
+
+    it('rejects when server is not in ready state', async () => {
+      server = new StudioBridgeServer({ placePath: '/fake/place.rbxl' });
+
+      await expect(
+        server.performActionAsync({ type: 'queryState', payload: {} })
+      ).rejects.toThrow(
+        "Cannot perform action: expected state 'ready', got 'idle'"
+      );
+    });
+  });
+
+  describe('health endpoint', () => {
+    it('GET /health returns 200 with correct JSON shape', async () => {
+      const ready = await createReadyServer({ sessionId: 'health-session' });
+      server = ready.server;
+      client = ready.client;
+
+      const res = await fetch(`http://localhost:${ready.port}/health`);
+      expect(res.status).toBe(200);
+      expect(res.headers.get('content-type')).toBe('application/json');
+
+      const body = (await res.json()) as Record<string, unknown>;
+      expect(body.status).toBe('ok');
+      expect(body.sessionId).toBe('health-session');
+      expect(body.port).toBe(ready.port);
+      expect(typeof body.serverVersion).toBe('string');
+    });
+
+    it('GET /other returns 404', async () => {
+      const ready = await createReadyServer();
+      server = ready.server;
+      client = ready.client;
+
+      const res = await fetch(`http://localhost:${ready.port}/other`);
+      expect(res.status).toBe(404);
+
+      const body = await res.text();
+      expect(body).toBe('Not Found');
+    });
+
+    it('health endpoint is available immediately after startAsync', async () => {
+      const sessionId = 'immediate-health';
+      server = new StudioBridgeServer({
+        placePath: '/fake/place.rbxl',
+        sessionId,
+        timeoutMs: 5_000,
+      });
+
+      const startPromise = server.startAsync();
+      await new Promise((r) => setTimeout(r, 50));
+      const port: number = (server as any)._port;
+
+      // Health should be available even before handshake completes
+      const res = await fetch(`http://localhost:${port}/health`);
+      expect(res.status).toBe(200);
+
+      const body = (await res.json()) as Record<string, unknown>;
+      expect(body.status).toBe('ok');
+      expect(body.sessionId).toBe(sessionId);
+
+      // Complete the handshake so startAsync resolves
+      client = await connectAndHandshake(port, sessionId);
+      await startPromise;
+    });
+
+    it('WebSocket upgrades to correct path still work after adding health endpoint', async () => {
+      const ready = await createReadyServer({ sessionId: 'ws-with-health' });
+      server = ready.server;
+      client = ready.client;
+
+      // If we got here, the WebSocket handshake succeeded through the HTTP server
+      // Verify health also works
+      const res = await fetch(`http://localhost:${ready.port}/health`);
+      expect(res.status).toBe(200);
+    });
+  });
 });