@kb-labs/gateway-app 0.2.0

package/src/__tests__/ws-client-connect.e2e.test.ts

@@ -0,0 +1,381 @@
+/**
+ * E2E WebSocket tests for /clients/connect (CC5 — Multi-Client Pub/Sub).
+ * Spins up a real Fastify server on a random port.
+ *
+ * Tests the full client protocol:
+ *   auth → client:hello → client:connected → subscribe/unsubscribe/cancel
+ *
+ * Does NOT mix with host-side ws-handler tests.
+ */
+import { describe, it, expect, beforeAll, afterAll, beforeEach } from 'vitest';
+import Fastify, { type FastifyInstance } from 'fastify';
+import fastifyWebsocket from '@fastify/websocket';
+import fastifyCors from '@fastify/cors';
+import { WebSocket, type RawData } from 'ws';
+import type { ICache, ILogger } from '@kb-labs/core-platform';
+import type { JwtConfig } from '@kb-labs/gateway-auth';
+import { AuthService } from '@kb-labs/gateway-auth';
+import { createClientWsHandler } from '../clients/ws-handler.js';
+import { executionRegistry } from '../execute/execution-registry.js';
+
+// ── In-memory ICache ─────────────────────────────────────────────────────────
+
+function makeInMemoryCache(): ICache {
+  const store = new Map<string, unknown>();
+  return {
+    async get<T>(key: string): Promise<T | null> {
+      return (store.get(key) as T) ?? null;
+    },
+    async set(key: string, value: unknown): Promise<void> {
+      store.set(key, value);
+    },
+    async delete(key: string): Promise<void> {
+      store.delete(key);
+    },
+    async clear(): Promise<void> {
+      store.clear();
+    },
+  } as unknown as ICache;
+}
+
+// ── Minimal ILogger ──────────────────────────────────────────────────────────
+
+const noopLogger: ILogger = {
+  info: () => {},
+  warn: () => {},
+  error: () => {},
+  debug: () => {},
+  child: () => noopLogger,
+} as unknown as ILogger;
+
+// ── Server setup ─────────────────────────────────────────────────────────────
+
+const testJwtConfig: JwtConfig = { secret: 'test-client-secret' };
+
+let app: FastifyInstance;
+let wsUrl: string;
+let cache: ICache;
+let authService: AuthService;
+
+beforeAll(async () => {
+  cache = makeInMemoryCache();
+  authService = new AuthService(cache, testJwtConfig);
+
+  app = Fastify({ logger: false });
+  await app.register(fastifyWebsocket);
+  await app.register(fastifyCors, { origin: true });
+
+  app.get('/clients/connect', { websocket: true }, createClientWsHandler(cache, testJwtConfig, noopLogger));
+
+  const address = await app.listen({ port: 0, host: '127.0.0.1' });
+  wsUrl = address.replace('http://', 'ws://');
+}, 10_000);
+
+afterAll(async () => {
+  await app.close();
+}, 10_000);
+
+beforeEach(() => {
+  // Clean up any leftover executions
+  for (const id of ['exec-ws-test-1', 'exec-ws-test-2', 'exec-ns-other']) {
+    executionRegistry.remove(id);
+  }
+});
+
+// ── Helpers ──────────────────────────────────────────────────────────────────
+
+/** Get a valid access token via AuthService.register + issueTokens */
+async function getAccessToken(namespaceId = 'ns-client-e2e'): Promise<string> {
+  const reg = await authService.register({ name: 'test-client', namespaceId, capabilities: [] });
+  const tokens = await authService.issueTokens(reg.clientId, reg.clientSecret);
+  return tokens!.accessToken;
+}
+
+/** Connect to /clients/connect with Bearer token */
+function connectClient(token: string): WebSocket {
+  return new WebSocket(`${wsUrl}/clients/connect`, {
+    headers: { Authorization: `Bearer ${token}` },
+  });
+}
+
+/** Collect the next N messages from a WS connection */
+function collectMessages(ws: WebSocket, count: number, timeout = 3000): Promise<unknown[]> {
+  return new Promise((resolve, reject) => {
+    const msgs: unknown[] = [];
+    const timer = setTimeout(() => {
+      reject(new Error(`Timeout: expected ${count} messages, got ${msgs.length}: ${JSON.stringify(msgs)}`));
+    }, timeout);
+
+    ws.on('message', (raw: RawData) => {
+      msgs.push(JSON.parse(raw.toString()));
+      if (msgs.length >= count) {
+        clearTimeout(timer);
+        resolve(msgs);
+      }
+    });
+
+    ws.on('error', (err: Error) => { clearTimeout(timer); reject(err); });
+  });
+}
+
+/** Wait for WS to open */
+function waitForOpen(ws: WebSocket): Promise<void> {
+  return new Promise((resolve, reject) => {
+    ws.on('open', resolve);
+    ws.on('error', reject);
+  });
+}
+
+// ── Tests ────────────────────────────────────────────────────────────────────
+
+describe('WebSocket /clients/connect: auth rejection', () => {
+  it('closes immediately when no Authorization header', async () => {
+    const ws = new WebSocket(`${wsUrl}/clients/connect`); // no auth
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('close', (code: number) => resolve(code));
+      ws.on('error', () => {});
+      setTimeout(() => resolve(-1), 3000);
+    });
+    expect([1008, 1006]).toContain(closeCode);
+  }, 5000);
+
+  it('closes with 1008 for invalid token', async () => {
+    const ws = new WebSocket(`${wsUrl}/clients/connect`, {
+      headers: { Authorization: 'Bearer not-a-valid-token' },
+    });
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('close', (code: number) => resolve(code));
+      ws.on('error', () => {});
+      setTimeout(() => resolve(-1), 3000);
+    });
+    expect([1008, 1006]).toContain(closeCode);
+  }, 5000);
+});
+
+describe('WebSocket /clients/connect: handshake', () => {
+  it('completes client:hello → client:connected flow', async () => {
+    const token = await getAccessToken();
+    const ws = connectClient(token);
+
+    await waitForOpen(ws);
+
+    const messagesP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:hello', clientVersion: '0.1.0' }));
+
+    const [connected] = await messagesP as [{ type: string; protocolVersion: string; connectionId: string }];
+    expect(connected.type).toBe('client:connected');
+    expect(connected.protocolVersion).toBe('1.0');
+    expect(typeof connected.connectionId).toBe('string');
+
+    ws.close(1000);
+  }, 8000);
+
+  it('closes with 1008 if client:hello not sent within timeout', async () => {
+    const token = await getAccessToken();
+    const ws = connectClient(token);
+
+    await waitForOpen(ws);
+    // Don't send hello — wait for timeout
+
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('close', (code: number) => resolve(code));
+      setTimeout(() => resolve(-1), 8000);
+    });
+
+    expect(closeCode).toBe(1008);
+  }, 10_000);
+
+  it('closes with 1008 if client:hello has invalid format', async () => {
+    const token = await getAccessToken();
+    const ws = connectClient(token);
+
+    await waitForOpen(ws);
+    ws.send(JSON.stringify({ type: 'client:hello' })); // missing clientVersion
+
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('close', (code: number) => resolve(code));
+      ws.on('error', () => {});
+      setTimeout(() => resolve(-1), 3000);
+    });
+
+    expect([1008, 1000, 1005]).toContain(closeCode);
+  }, 5000);
+
+  it('access_token query param is accepted as fallback auth', async () => {
+    const token = await getAccessToken();
+    const ws = new WebSocket(`${wsUrl}/clients/connect?access_token=${token}`);
+
+    await waitForOpen(ws);
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:hello', clientVersion: '0.1.0' }));
+
+    const [connected] = await msgsP as [{ type: string }];
+    expect(connected.type).toBe('client:connected');
+
+    ws.close(1000);
+  }, 8000);
+});
+
+describe('WebSocket /clients/connect: subscribe / unsubscribe', () => {
+  async function connectAndHandshake(): Promise<{ ws: WebSocket; connectionId: string }> {
+    const token = await getAccessToken();
+    const ws = connectClient(token);
+    await waitForOpen(ws);
+
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:hello', clientVersion: '0.1.0' }));
+    const [connected] = await msgsP as [{ type: string; connectionId: string }];
+    return { ws, connectionId: connected.connectionId };
+  }
+
+  it('receives client:error(EXECUTION_NOT_FOUND) when subscribing to unknown execution', async () => {
+    const { ws } = await connectAndHandshake();
+
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:subscribe', executionId: '00000000-0000-0000-0000-000000000001' }));
+
+    const [errorMsg] = await msgsP as [{ type: string; code: string }];
+    expect(errorMsg.type).toBe('client:error');
+    expect(errorMsg.code).toBe('EXECUTION_NOT_FOUND');
+
+    ws.close(1000);
+  }, 8000);
+
+  it('receives client:error(FORBIDDEN) when subscribing to execution in different namespace', async () => {
+    // Register an execution in a different namespace
+    const execId = '00000000-0000-0000-0000-000000000002';
+    executionRegistry.register({
+      executionId: execId,
+      requestId: 'req-x',
+      namespaceId: 'ns-other-forbidden',  // different namespace
+      hostId: 'host-x',
+      pluginId: 'p',
+      handlerRef: 'h',
+    });
+
+    const { ws } = await connectAndHandshake();
+
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:subscribe', executionId: execId }));
+
+    const [errorMsg] = await msgsP as [{ type: string; code: string }];
+    expect(errorMsg.type).toBe('client:error');
+    expect(errorMsg.code).toBe('FORBIDDEN');
+
+    ws.close(1000);
+    executionRegistry.remove(execId);
+  }, 8000);
+
+  it('does not error on unsubscribe from unknown execution (graceful)', async () => {
+    const { ws } = await connectAndHandshake();
+
+    // Send unsubscribe without subscribing first — should not crash
+    ws.send(JSON.stringify({ type: 'client:unsubscribe', executionId: '00000000-0000-0000-0000-000000000003' }));
+
+    // If we get here without the connection closing, it's fine
+    await new Promise((r) => { setTimeout(r, 200); });
+    expect(ws.readyState).toBe(ws.OPEN);
+
+    ws.close(1000);
+  }, 5000);
+
+  it('receives client:error(INVALID_MESSAGE) for malformed JSON', async () => {
+    const { ws } = await connectAndHandshake();
+
+    const msgsP = collectMessages(ws, 1);
+    ws.send('not json at all {{{');
+
+    const [errorMsg] = await msgsP as [{ type: string; code: string }];
+    expect(errorMsg.type).toBe('client:error');
+    expect(errorMsg.code).toBe('INVALID_MESSAGE');
+
+    ws.close(1000);
+  }, 5000);
+
+  it('receives client:error(INVALID_MESSAGE) for unknown message type', async () => {
+    const { ws } = await connectAndHandshake();
+
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'unknown:type' }));
+
+    const [errorMsg] = await msgsP as [{ type: string; code: string }];
+    expect(errorMsg.type).toBe('client:error');
+    expect(errorMsg.code).toBe('INVALID_MESSAGE');
+
+    ws.close(1000);
+  }, 5000);
+});
+
+describe('WebSocket /clients/connect: cancel', () => {
+  async function connectAndHandshake(): Promise<WebSocket> {
+    const token = await getAccessToken();
+    const ws = connectClient(token);
+    await waitForOpen(ws);
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:hello', clientVersion: '0.1.0' }));
+    await msgsP;
+    return ws;
+  }
+
+  it('receives client:error(EXECUTION_NOT_FOUND) when cancelling unknown execution', async () => {
+    const ws = await connectAndHandshake();
+
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:cancel', executionId: '00000000-0000-0000-0000-000000000004' }));
+
+    const [errorMsg] = await msgsP as [{ type: string; code: string }];
+    expect(errorMsg.type).toBe('client:error');
+    expect(errorMsg.code).toBe('EXECUTION_NOT_FOUND');
+
+    ws.close(1000);
+  }, 8000);
+
+  it('cancels execution and aborts the signal', async () => {
+    const execId = '00000000-0000-0000-0000-000000000005';
+    const signal = executionRegistry.register({
+      executionId: execId,
+      requestId: 'req-cancel',
+      namespaceId: 'ns-client-e2e', // same namespace as token
+      hostId: 'host-001',
+      pluginId: 'p',
+      handlerRef: 'h',
+    });
+
+    const ws = await connectAndHandshake();
+    ws.send(JSON.stringify({ type: 'client:cancel', executionId: execId, reason: 'user' }));
+
+    // Give the server time to process
+    await new Promise((r) => { setTimeout(r, 100); });
+
+    expect(signal.aborted).toBe(true);
+
+    ws.close(1000);
+    executionRegistry.remove(execId);
+  }, 8000);
+
+  it('receives client:error(CANCEL_FAILED) when execution already cancelled', async () => {
+    const execId = '00000000-0000-0000-0000-000000000006';
+    executionRegistry.register({
+      executionId: execId,
+      requestId: 'req-already',
+      namespaceId: 'ns-client-e2e',
+      hostId: 'host-001',
+      pluginId: 'p',
+      handlerRef: 'h',
+    });
+    // Cancel it first
+    executionRegistry.cancel(execId, 'user');
+
+    const ws = await connectAndHandshake();
+
+    const msgsP = collectMessages(ws, 1);
+    ws.send(JSON.stringify({ type: 'client:cancel', executionId: execId, reason: 'user' }));
+
+    const [errorMsg] = await msgsP as [{ type: string; code: string }];
+    expect(errorMsg.type).toBe('client:error');
+    expect(errorMsg.code).toBe('CANCEL_FAILED');
+
+    ws.close(1000);
+    executionRegistry.remove(execId);
+  }, 8000);
+});

package/src/__tests__/ws-handshake.e2e.test.ts

@@ -0,0 +1,288 @@
+/**
+ * E2E WebSocket tests — spins up a real Fastify server on a random port.
+ * Tests the full handshake path: auth → hello → connected → heartbeat → ack → close.
+ */
+import { describe, it, expect, beforeAll, afterAll } from 'vitest';
+import Fastify, { type FastifyInstance } from 'fastify';
+import fastifyWebsocket from '@fastify/websocket';
+import fastifyCors from '@fastify/cors';
+import { WebSocket, type RawData } from 'ws';
+import type { ICache } from '@kb-labs/core-platform';
+import { HostRegistrationSchema } from '@kb-labs/gateway-contracts';
+import type { JwtConfig } from '@kb-labs/gateway-auth';
+import { createAuthMiddleware } from '../auth/middleware.js';
+import { HostRegistry } from '../hosts/registry.js';
+import { createWsHandler } from '../hosts/ws-handler.js';
+import type { ILogger } from '@kb-labs/core-platform';
+
+// Minimal jwtConfig for tests — JWT auth is disabled (empty secret forces fallback to cache tokens)
+const testJwtConfig: JwtConfig = { secret: 'test-secret-do-not-use' };
+
+// ── In-memory ICache ─────────────────────────────────────────────────────────
+
+function makeInMemoryCache(): ICache {
+  const store = new Map<string, unknown>();
+  return {
+    async get<T>(key: string): Promise<T | null> {
+      return (store.get(key) as T) ?? null;
+    },
+    async set(key: string, value: unknown): Promise<void> {
+      store.set(key, value);
+    },
+    async delete(key: string): Promise<void> {
+      store.delete(key);
+    },
+    async clear(): Promise<void> {
+      store.clear();
+    },
+  } as unknown as ICache;
+}
+
+// ── Server setup ─────────────────────────────────────────────────────────────
+
+let app: FastifyInstance;
+let baseUrl: string;
+let wsUrl: string;
+let cache: ICache;
+const noopLogger: ILogger = {
+  info: () => {},
+  warn: () => {},
+  error: () => {},
+  debug: () => {},
+  child: () => noopLogger,
+} as unknown as ILogger;
+
+beforeAll(async () => {
+  cache = makeInMemoryCache();
+  app = Fastify({ logger: false });
+
+  await app.register(fastifyWebsocket);
+  await app.register(fastifyCors, { origin: true });
+  app.addHook('preHandler', createAuthMiddleware(cache, testJwtConfig));
+
+  app.get('/health', async () => ({ status: 'ok', version: '1.0' }));
+
+  const registry = new HostRegistry(cache);
+  app.post('/hosts/register', async (request, reply) => {
+    const parsed = HostRegistrationSchema.safeParse(request.body);
+    if (!parsed.success) {return reply.code(400).send({ error: 'Bad Request' });}
+    const result = await registry.register(parsed.data);
+    return reply.code(201).send({
+      hostId: result.descriptor.hostId,
+      machineToken: result.machineToken,
+      status: result.descriptor.status,
+    });
+  });
+
+  app.get('/hosts/connect', { websocket: true }, createWsHandler(cache, testJwtConfig, noopLogger));
+
+  const address = await app.listen({ port: 0, host: '127.0.0.1' }); // port 0 = random
+  baseUrl = address;
+  wsUrl = address.replace('http://', 'ws://');
+}, 10_000);
+
+afterAll(async () => {
+  await app.close();
+}, 10_000);
+
+// ── Helpers ──────────────────────────────────────────────────────────────────
+
+async function registerHost(name = 'test-host'): Promise<{ hostId: string; machineToken: string }> {
+  const res = await fetch(`${baseUrl}/hosts/register`, {
+    method: 'POST',
+    headers: { 'content-type': 'application/json' },
+    body: JSON.stringify({ name, namespaceId: 'ns-e2e', capabilities: ['filesystem'], workspacePaths: [] }),
+  });
+  return res.json() as Promise<{ hostId: string; machineToken: string }>;
+}
+
+function connectWs(token: string): WebSocket {
+  return new WebSocket(`${wsUrl}/hosts/connect`, {
+    headers: { Authorization: `Bearer ${token}` },
+  });
+}
+
+function collectMessages(ws: WebSocket, count: number, timeout = 3000): Promise<unknown[]> {
+  return new Promise((resolve, reject) => {
+    const msgs: unknown[] = [];
+    const timer = setTimeout(() => {
+      reject(new Error(`Timeout: expected ${count} messages, got ${msgs.length}: ${JSON.stringify(msgs)}`));
+    }, timeout);
+
+    ws.on('message', (raw: RawData) => {
+      msgs.push(JSON.parse(raw.toString()));
+      if (msgs.length >= count) {
+        clearTimeout(timer);
+        resolve(msgs);
+      }
+    });
+
+    ws.on('error', (err: Error) => { clearTimeout(timer); reject(err); });
+  });
+}
+
+// ── Tests ────────────────────────────────────────────────────────────────────
+
+describe('WebSocket: connection refused without auth', () => {
+  it('closes immediately when no Authorization header', async () => {
+    const ws = new WebSocket(`${wsUrl}/hosts/connect`); // no auth header
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('close', (code: number) => resolve(code));
+      ws.on('error', () => {}); // suppress
+    });
+    // 1008 = Policy Violation (explicit close from server)
+    // 1006 = Abnormal closure (server drops connection at transport level)
+    // Both indicate the connection was rejected — either is correct depending on Fastify/ws version
+    expect([1008, 1006]).toContain(closeCode);
+  }, 5000);
+
+  it('closes with 1008 for invalid machine token', async () => {
+    const ws = connectWs('not-a-valid-machine-token');
+    // ws token resolves as CLI (not machine) — should be rejected in ws-handler
+    // The handler checks tokenEntry.type !== 'machine'
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('close', (code: number) => resolve(code));
+      ws.on('open', () => {
+        // If opened, send hello — handler will auth-check and close
+      });
+      ws.on('error', () => resolve(1008));
+      setTimeout(() => resolve(1008), 3000);
+    });
+    // Either 1008 or connection never established
+    expect([1008, 1000, 1005]).toContain(closeCode);
+  }, 5000);
+});
+
+describe('WebSocket: full handshake', () => {
+  it('completes hello → connected flow', async () => {
+    const { machineToken, hostId } = await registerHost('e2e-host-1');
+    const ws = connectWs(machineToken);
+
+    const messages = collectMessages(ws, 1);
+
+    await new Promise<void>((resolve, reject) => {
+      ws.on('open', () => {
+        ws.send(JSON.stringify({ type: 'hello', protocolVersion: '1.0', agentVersion: '0.1.0' }));
+        resolve();
+      });
+      ws.on('error', reject);
+    });
+
+    const [connected] = await messages as [{ type: string; hostId: string; sessionId: string; protocolVersion: string }];
+    expect(connected.type).toBe('connected');
+    expect(connected.hostId).toBe(hostId);
+    expect(connected.protocolVersion).toBe('1.0');
+    expect(connected.sessionId).toBeTypeOf('string');
+
+    ws.close(1000);
+  }, 8000);
+
+  it('heartbeat gets ack response', async () => {
+    const { machineToken } = await registerHost('e2e-host-2');
+    const ws = connectWs(machineToken);
+
+    await new Promise<void>((resolve, reject) => {
+      ws.on('open', () => {
+        ws.send(JSON.stringify({ type: 'hello', protocolVersion: '1.0', agentVersion: '0.1.0' }));
+        resolve();
+      });
+      ws.on('error', reject);
+    });
+
+    // Wait for 'connected', then send heartbeat
+    const ack = await new Promise<unknown>((resolve, reject) => {
+      let gotConnected = false;
+      const timer = setTimeout(() => reject(new Error('timeout')), 5000);
+
+      ws.on('message', (raw: RawData) => {
+        const msg = JSON.parse(raw.toString()) as { type: string };
+        if (msg.type === 'connected' && !gotConnected) {
+          gotConnected = true;
+          ws.send(JSON.stringify({ type: 'heartbeat' }));
+        } else if (msg.type === 'ack') {
+          clearTimeout(timer);
+          resolve(msg);
+        }
+      });
+      ws.on('error', (e) => { clearTimeout(timer); reject(e); });
+    });
+
+    expect((ack as { type: string }).type).toBe('ack');
+    ws.close(1000);
+  }, 8000);
+
+  it('rejects unsupported protocol version', async () => {
+    const { machineToken } = await registerHost('e2e-host-3');
+    const ws = connectWs(machineToken);
+
+    const msgs: unknown[] = [];
+    await new Promise<void>((resolve, reject) => {
+      ws.on('open', () => {
+        ws.send(JSON.stringify({ type: 'hello', protocolVersion: '99.0', agentVersion: '0.1.0' }));
+        resolve();
+      });
+      ws.on('error', reject);
+    });
+
+    // Should receive 'negotiate' message and then close
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('message', (raw: RawData) => { msgs.push(JSON.parse(raw.toString())); });
+      ws.on('close', (code: number) => resolve(code));
+      setTimeout(() => resolve(1008), 4000);
+    });
+
+    expect(closeCode).toBe(1008);
+    expect((msgs[0] as { type: string; supportedVersions: string[] }).type).toBe('negotiate');
+    expect((msgs[0] as { type: string; supportedVersions: string[] }).supportedVersions).toContain('1.0');
+  }, 8000);
+
+  it('closes with 1008 if hello not sent within timeout', async () => {
+    const { machineToken } = await registerHost('e2e-host-4');
+    const ws = connectWs(machineToken);
+
+    // Open but don't send hello
+    await new Promise<void>((resolve, reject) => {
+      ws.on('open', () => resolve());
+      ws.on('error', reject);
+    });
+
+    const closeCode = await new Promise<number>((resolve) => {
+      ws.on('close', (code: number) => resolve(code));
+      setTimeout(() => resolve(-1), 8000); // wait longer than HELLO_TIMEOUT_MS (5s)
+    });
+
+    expect(closeCode).toBe(1008);
+  }, 10_000);
+});
+
+describe('WebSocket: host status lifecycle', () => {
+  it('host goes online after successful handshake', async () => {
+    const { machineToken, hostId } = await registerHost('e2e-status-1');
+
+    const ws = connectWs(machineToken);
+    await new Promise<void>((resolve, reject) => {
+      ws.on('open', () => {
+        ws.send(JSON.stringify({ type: 'hello', protocolVersion: '1.0', agentVersion: '0.1.0' }));
+      });
+      ws.on('message', (raw: RawData) => {
+        const msg = JSON.parse(raw.toString()) as { type: string };
+        if (msg.type === 'connected') {resolve();}
+      });
+      ws.on('error', reject);
+    });
+
+    // Verify host is online in cache
+    const host = await cache.get<{ status: string; connections: string[] }>(`host:registry:ns-e2e:${hostId}`);
+    expect(host!.status).toBe('online');
+    expect(host!.connections.length).toBeGreaterThan(0);
+
+    ws.close(1000);
+
+    // Give server time to process close — status transitions to 'reconnecting' first,
+    // then to 'offline' after the reconnect grace period
+    await new Promise((r) => { setTimeout(r, 100); });
+
+    const hostAfter = await cache.get<{ status: string }>(`host:registry:ns-e2e:${hostId}`);
+    expect(['reconnecting', 'offline']).toContain(hostAfter!.status);
+  }, 8000);
+});