test(js/client): cover disconnect-drain for Publication and Connection

Three node:test cases against a real local ws.WebSocketServer (no
mocks): ZhubPublication.invoke() rejects fast on remote terminate;
ZhubConnection.chat() rejects fast on remote terminate; and a happy-
path chat() echo so the rejection tests aren't the only coverage of
the chat path. Both rejection cases assert ZhubConnectionError and
elapsed < 2s with a 60s timeoutMs configured, so a regression here
manifests as a multi-second hang.

Add the new client.test.js bundle to the npm test target.

Mutation-verified: reverting js/src/client.ts leaves the elapsed
assertion failing (call hangs the full 60s before the per-call
timeout fires).

Author: Zawwarsami16

js/package.json

@@ -26,7 +26,7 @@
   ],
   "scripts": {
     "build": "tsc -p tsconfig.json",
-    "test": "tsc -p tsconfig.test.json && node --test dist-test/test/manifest.test.js dist-test/test/protocol.test.js"
+    "test": "tsc -p tsconfig.test.json && node --test dist-test/test/manifest.test.js dist-test/test/protocol.test.js dist-test/test/client.test.js"
   },
   "dependencies": {
     "ws": "^8.18.0"

js/test/client.test.ts

@@ -0,0 +1,165 @@
+/**
+ * ZhubPublication / ZhubConnection disconnect-drain tests.
+ *
+ * Regression for: pending invoke()/chat() futures waited the full 60-second
+ * timeout when the WS dropped instead of rejecting immediately on close.
+ *
+ * Uses a real local ws.WebSocketServer so the transport is faithful.
+ */
+import { describe, it, before, after } from 'node:test';
+import assert from 'node:assert/strict';
+import { WebSocketServer } from 'ws';
+import type { AddressInfo } from 'node:net';
+import { ZhubPublication, ZhubConnection, publish, connect } from '../src/client.js';
+import { ZhubConnectionError } from '../src/errors.js';
+import { chatResponse } from '../src/protocol.js';
+
+function makeHubUrl(port: number): string {
+  return `http://127.0.0.1:${port}`;
+}
+
+async function startServer(): Promise<{ wss: WebSocketServer; port: number; close: () => Promise<void> }> {
+  const wss = new WebSocketServer({ port: 0 });
+  await new Promise<void>((r) => wss.on('listening', r));
+  const port = (wss.address() as AddressInfo).port;
+  return {
+    wss,
+    port,
+    close: () => new Promise<void>((r) => wss.close(() => r())),
+  };
+}
+
+describe('ZhubPublication disconnect drain', () => {
+  let wss: WebSocketServer;
+  let port: number;
+  let closeServer: () => Promise<void>;
+
+  before(async () => {
+    const s = await startServer();
+    wss = s.wss;
+    port = s.port;
+    closeServer = s.close;
+
+    // Accept connections and send 'registered', then stay open so callers can invoke
+    wss.on('connection', (ws) => {
+      ws.send(JSON.stringify({ type: 'registered', request_id: 'r0', payload: { name: 'test-pub', base_url: '', api_key: 'zk_test' } }));
+    });
+  });
+
+  after(async () => {
+    await closeServer();
+  });
+
+  it('invoke() rejects with ZhubConnectionError immediately when WS closes', async () => {
+    const pub = publish({
+      name: 'test-pub',
+      description: 'test',
+      hubUrl: makeHubUrl(port),
+      apiKey: 'zk_test',
+      chatHandler: async () => 'hi',
+    });
+
+    // Wait until the WS is open (registered message received)
+    await new Promise<void>((r) => setTimeout(r, 100));
+
+    // Kick off an invoke() — no server-side handler, so it would wait for the timeout
+    const invokePromise = pub.invoke('cx_1', 'ping', {}, 60_000);
+
+    // Kill all server-side connections to trigger onclose on the client
+    for (const client of wss.clients) client.terminate();
+
+    const start = Date.now();
+    await assert.rejects(invokePromise, ZhubConnectionError);
+    const elapsed = Date.now() - start;
+    assert(elapsed < 2_000, `expected fast rejection, got ${elapsed}ms`);
+
+    await pub.stop();
+  });
+});
+
+describe('ZhubConnection disconnect drain', () => {
+  let wss: WebSocketServer;
+  let port: number;
+  let closeServer: () => Promise<void>;
+
+  before(async () => {
+    const s = await startServer();
+    wss = s.wss;
+    port = s.port;
+    closeServer = s.close;
+
+    wss.on('connection', (ws) => {
+      ws.send(JSON.stringify({ type: 'registered', request_id: 'r0', payload: {} }));
+    });
+  });
+
+  after(async () => {
+    await closeServer();
+  });
+
+  it('chat() rejects with ZhubConnectionError immediately when WS closes', async () => {
+    const conn = connect({
+      aiName: 'test-ai',
+      apiKey: 'zk_test',
+      hubUrl: makeHubUrl(port),
+    });
+
+    // Wait until registered
+    await new Promise<void>((r) => setTimeout(r, 100));
+
+    const chatPromise = conn.chat([{ role: 'user', content: 'hello' }], { timeoutMs: 60_000 });
+
+    // Close all server-side sockets
+    for (const client of wss.clients) client.terminate();
+
+    const start = Date.now();
+    await assert.rejects(chatPromise, ZhubConnectionError);
+    const elapsed = Date.now() - start;
+    assert(elapsed < 2_000, `expected fast rejection, got ${elapsed}ms`);
+
+    await conn.stop();
+  });
+});
+
+describe('ZhubConnection chat() happy path', () => {
+  let wss: WebSocketServer;
+  let port: number;
+  let closeServer: () => Promise<void>;
+
+  before(async () => {
+    const s = await startServer();
+    wss = s.wss;
+    port = s.port;
+    closeServer = s.close;
+
+    // Echo server: on chat-request, immediately send chat-response
+    wss.on('connection', (ws) => {
+      ws.send(JSON.stringify({ type: 'registered', request_id: 'r0', payload: {} }));
+      ws.on('message', (raw) => {
+        const env = JSON.parse(raw.toString());
+        if (env.type === 'chat-request') {
+          ws.send(JSON.stringify(chatResponse('hello back', env.request_id, 'stop')));
+        }
+      });
+    });
+  });
+
+  after(async () => {
+    await closeServer();
+  });
+
+  it('resolves with text from hub response', async () => {
+    const conn = connect({
+      aiName: 'test-ai',
+      apiKey: 'zk_test',
+      hubUrl: makeHubUrl(port),
+    });
+
+    await new Promise<void>((r) => setTimeout(r, 100));
+
+    const result = await conn.chat([{ role: 'user', content: 'hello' }], { timeoutMs: 5_000 });
+    assert.equal(result.text, 'hello back');
+
+    await conn.stop();
+  });
+});