@kortix/agent-tunnel 0.1.4 → 0.12.7

package/src/client/tunnel-client.ts

@@ -5,6 +5,12 @@ export interface TunnelClientConfig {
   cacheTtlMs?: number;
 }
 
+function trimTrailingSlashes(value: string): string {
+  let end = value.length;
+  while (end > 0 && value.charCodeAt(end - 1) === 47) end -= 1;
+  return value.slice(0, end);
+}
+
 export class TunnelClientError extends Error {
   constructor(
     public readonly code: number,
@@ -27,17 +33,26 @@ export class TunnelClient {
 
   readonly fs: FsNamespace;
   readonly shell: ShellNamespace;
-  readonly desktop: DesktopNamespace;
+  readonly cua: CuaNamespace;
 
   constructor(config: TunnelClientConfig) {
-    this.apiUrl = config.apiUrl.replace(/\/+$/, '');
+    const apiUrl = new URL(config.apiUrl);
+    const loopback =
+      apiUrl.hostname === 'localhost' ||
+      apiUrl.hostname === '127.0.0.1' ||
+      apiUrl.hostname === '[::1]' ||
+      apiUrl.hostname === '::1';
+    if (apiUrl.protocol !== 'https:' && !(apiUrl.protocol === 'http:' && loopback)) {
+      throw new Error('Remote tunnel API URLs must use https');
+    }
+    this.apiUrl = trimTrailingSlashes(apiUrl.toString());
     this.token = config.token;
     this.explicitTunnelId = config.tunnelId;
     this.cacheTtlMs = config.cacheTtlMs ?? 10_000;
 
     this.fs = new FsNamespace(this);
     this.shell = new ShellNamespace(this);
-    this.desktop = new DesktopNamespace(this);
+    this.cua = new CuaNamespace(this);
   }
 
   async rpc(method: string, params: Record<string, unknown> = {}): Promise<unknown> {
@@ -109,15 +124,11 @@ export class TunnelClient {
     if (res.ok) {
       const connections = (await res.json()) as Array<{ tunnelId: string; isLive?: boolean }>;
       const online = connections.find((c) => c.isLive);
-      if (online) {
-        this.cachedTunnelId = online.tunnelId;
+      const chosen = online ?? connections[0];
+      if (chosen) {
+        this.cachedTunnelId = chosen.tunnelId;
         this.cacheTimestamp = Date.now();
-        return online.tunnelId;
-      }
-      if (connections.length > 0) {
-        this.cachedTunnelId = connections[0].tunnelId;
-        this.cacheTimestamp = Date.now();
-        return connections[0].tunnelId;
+        return chosen.tunnelId;
       }
     }
 
@@ -126,7 +137,7 @@ export class TunnelClient {
       -1,
       'No tunnel connection found. The user needs to set up Agent Tunnel first:\n' +
       '1. Create a tunnel connection\n' +
-      '2. Run `npx @kortix/agent-tunnel connect` on their local machine',
+      '2. Connect the local machine from the Kortix desktop app or run the tunnel connect command',
     );
   }
 }
@@ -172,99 +183,58 @@ class ShellNamespace {
   }
 }
 
-class DesktopNamespace {
+class CuaNamespace {
   constructor(private client: TunnelClient) {}
 
-  async screenshot(params?: { region?: { x: number; y: number; width: number; height: number }; windowId?: number }): Promise<{ image: string; width: number; height: number; format?: string }> {
-    return (await this.client.rpc('desktop.screenshot', params ?? {})) as any;
-  }
-
-  async click(params: { x: number; y: number; button?: string; clicks?: number; modifiers?: string[] }): Promise<unknown> {
-    return this.client.rpc('desktop.mouse.click', params);
-  }
-
-  async type(text: string, delay?: number): Promise<unknown> {
-    return this.client.rpc('desktop.keyboard.type', { text, delay });
-  }
-
-  async key(keys: string[]): Promise<unknown> {
-    return this.client.rpc('desktop.keyboard.key', { keys });
+  async ensure(): Promise<{ ok: boolean; binary: string; version?: string }> {
+    return (await this.client.rpc('desktop.cua.ensure', {})) as any;
   }
 
-  async mouseMove(x: number, y: number): Promise<unknown> {
-    return this.client.rpc('desktop.mouse.move', { x, y });
+  async startDaemon(): Promise<Record<string, unknown>> {
+    return (await this.client.rpc('desktop.cua.start_daemon', {})) as any;
   }
 
-  async mouseDrag(fromX: number, fromY: number, toX: number, toY: number, button?: string): Promise<unknown> {
-    return this.client.rpc('desktop.mouse.drag', { fromX, fromY, toX, toY, button });
+  async status(): Promise<{ status: string }> {
+    return (await this.client.rpc('desktop.cua.status', {})) as any;
   }
 
-  async mouseScroll(x: number, y: number, deltaX?: number, deltaY?: number): Promise<unknown> {
-    return this.client.rpc('desktop.mouse.scroll', { x, y, deltaX, deltaY });
+  async version(): Promise<{ version: string }> {
+    return (await this.client.rpc('desktop.cua.version', {})) as any;
   }
 
-  async windowList(): Promise<{ windows: Array<{ id: number; app: string; title: string; bounds: { x: number; y: number; width: number; height: number }; minimized: boolean }> }> {
-    return (await this.client.rpc('desktop.window.list', {})) as any;
+  async listTools(): Promise<{ tools: string }> {
+    return (await this.client.rpc('desktop.cua.list_tools', {})) as any;
   }
 
-  async windowFocus(windowId: number): Promise<unknown> {
-    return this.client.rpc('desktop.window.focus', { windowId });
+  async describe(tool: string): Promise<{ description: string }> {
+    return (await this.client.rpc('desktop.cua.describe', { tool })) as any;
   }
 
-  async appLaunch(app: string): Promise<unknown> {
-    return this.client.rpc('desktop.app.launch', { app });
+  async call(tool: string, args: Record<string, unknown> = {}): Promise<unknown> {
+    return this.client.rpc('desktop.cua.call', { tool, args });
   }
 
-  async appQuit(app: string): Promise<unknown> {
-    return this.client.rpc('desktop.app.quit', { app });
+  async listApps(): Promise<unknown> {
+    return this.client.rpc('desktop.cua.list_apps', {});
   }
 
-  async clipboardRead(): Promise<{ text: string }> {
-    return (await this.client.rpc('desktop.clipboard.read', {})) as any;
+  async listWindows(params?: { pid?: number; on_screen_only?: boolean }): Promise<unknown> {
+    return this.client.rpc('desktop.cua.list_windows', params ?? {});
   }
 
-  async clipboardWrite(text: string): Promise<unknown> {
-    return this.client.rpc('desktop.clipboard.write', { text });
+  async getWindowState(params: { pid: number; window_id: number; query?: string; capture_mode?: 'som' | 'vision' | 'ax'; screenshot_out_file?: string; session?: string }): Promise<unknown> {
+    return this.client.rpc('desktop.cua.get_window_state', params);
   }
 
-  async screenInfo(): Promise<{ width: number; height: number; scaleFactor: number }> {
-    return (await this.client.rpc('desktop.screen.info', {})) as any;
+  async click(params: Record<string, unknown>): Promise<unknown> {
+    return this.client.rpc('desktop.cua.click', params);
   }
 
-  async cursorImage(radius?: number): Promise<{ image: string; width: number; height: number; format?: string }> {
-    return (await this.client.rpc('desktop.cursor.image', { radius })) as any;
+  async typeText(params: Record<string, unknown>): Promise<unknown> {
+    return this.client.rpc('desktop.cua.type_text', params);
   }
 
-  async axTree(params?: { pid?: number; maxDepth?: number; roles?: string[] }): Promise<{ root: AXElement; elementCount: number }> {
-    return (await this.client.rpc('desktop.ax.tree', params ?? {})) as any;
+  async hotkey(params: Record<string, unknown>): Promise<unknown> {
+    return this.client.rpc('desktop.cua.hotkey', params);
   }
-
-  async axAction(elementId: string, action: string, pid?: number): Promise<Record<string, unknown>> {
-    return (await this.client.rpc('desktop.ax.action', { elementId, action, pid })) as any;
-  }
-
-  async axSetValue(elementId: string, value: string, pid?: number): Promise<Record<string, unknown>> {
-    return (await this.client.rpc('desktop.ax.set_value', { elementId, value, pid })) as any;
-  }
-
-  async axFocus(elementId: string, pid?: number): Promise<Record<string, unknown>> {
-    return (await this.client.rpc('desktop.ax.focus', { elementId, pid })) as any;
-  }
-
-  async axSearch(query: string, params?: { role?: string; pid?: number; maxResults?: number }): Promise<{ elements: AXElement[] }> {
-    return (await this.client.rpc('desktop.ax.search', { query, ...params })) as any;
-  }
-}
-
-export interface AXElement {
-  id: string;
-  role: string;
-  title: string;
-  value: string;
-  description: string;
-  bounds: { x: number; y: number; width: number; height: number };
-  children: AXElement[];
-  actions: string[];
-  enabled: boolean;
-  focused: boolean;
 }

package/src/index.ts

@@ -8,6 +8,11 @@ export {
   timingSafeStringEqual,
   TunnelErrorCode,
   TunnelMethods,
+  capabilityForMethod,
+  desktopFeatureForMethod,
+  isTunnelCapability,
+  operationForMethod,
+  validateTunnelPermissionScope,
 } from './shared';
 
 export type {
@@ -33,6 +38,7 @@ export type {
   TunnelCapability,
   TunnelMethod,
   TunnelErrorCodeValue,
+  PermissionScopeValidationResult,
 } from './shared';
 
 // ─── Server: Relay ──────────────────────────────────────────────────────────
@@ -46,7 +52,7 @@ export type { TunnelServer } from './server';
 
 // ─── Client: SDK ────────────────────────────────────────────────────────────
 export { TunnelClient, TunnelClientError } from './client';
-export type { TunnelClientConfig, AXElement } from './client';
+export type { TunnelClientConfig } from './client';
 export { createTunnelTools } from './client';
 export type { TunnelToolDefinition, TunnelToolParameter } from './client';
 

package/src/node-ws-polyfill.test.ts

@@ -0,0 +1,18 @@
+import { describe, expect, test } from 'bun:test';
+import { pathToFileURL } from 'node:url';
+
+describe('node WebSocket polyfill', () => {
+  test('loads ws when the runtime has no global WebSocket', () => {
+    const polyfillUrl = pathToFileURL(`${import.meta.dir}/node-ws-polyfill.ts`).href;
+    const script = `
+      delete globalThis.WebSocket;
+      await import(${JSON.stringify(polyfillUrl)});
+      if (typeof globalThis.WebSocket !== 'function') process.exit(1);
+    `;
+
+    const result = Bun.spawnSync([process.execPath, '--eval', script]);
+
+    expect(result.exitCode).toBe(0);
+    expect(result.stderr.toString()).toBe('');
+  });
+});

package/src/node-ws-polyfill.ts

@@ -5,9 +5,9 @@
  */
 if (typeof globalThis.WebSocket === 'undefined') {
   try {
-    // eslint-disable-next-line @typescript-eslint/no-require-imports
-    const ws = require('ws');
-    globalThis.WebSocket = ws.default || ws;
+    const ws = await import('ws');
+    const WebSocketImpl = ws.WebSocket ?? ws.default;
+    globalThis.WebSocket = WebSocketImpl as unknown as typeof WebSocket;
   } catch {
     console.error(
       '[agent-tunnel] WebSocket is not available. Install the "ws" package or use Node.js 22+.',
@@ -15,3 +15,5 @@ if (typeof globalThis.WebSocket === 'undefined') {
     process.exit(1);
   }
 }
+
+export {};

package/src/server/heartbeat.ts

@@ -25,12 +25,11 @@ export class HeartbeatManager {
   start(): void {
     if (this.intervalHandle) return;
 
-    this.intervalHandle = setInterval(
-      () => this.tick(),
-      this.intervalMs,
-    );
+    this.intervalHandle = setInterval(() => this.tick(), this.intervalMs);
 
-    console.log(`[tunnel-heartbeat] Started (interval: ${this.intervalMs}ms, max missed: ${this.maxMissed})`);
+    console.log(
+      `[tunnel-heartbeat] Started (interval: ${this.intervalMs}ms, max missed: ${this.maxMissed})`,
+    );
   }
 
   stop(): void {
@@ -47,6 +46,11 @@ export class HeartbeatManager {
       missedPongs: 0,
       lastPongAt: Date.now(),
     });
+    // Collect the live handler list immediately. Waiting for the first interval
+    // can advertise stale DB capabilities for up to 30 seconds after connect.
+    this.relay.sendNotification(tunnelId, 'tunnel.ping', {
+      timestamp: Date.now(),
+    });
   }
 
   unregister(tunnelId: string): void {
@@ -75,8 +79,11 @@ export class HeartbeatManager {
       state.missedPongs++;
 
       if (state.missedPongs >= this.maxMissed) {
-        console.warn(`[tunnel-heartbeat] Agent ${tunnelId} missed ${state.missedPongs} pongs — timing out`);
+        console.warn(
+          `[tunnel-heartbeat] Agent ${tunnelId} missed ${state.missedPongs} pongs — timing out`,
+        );
         this.relay.emitEvent('agent:timeout', { tunnelId });
+        this.relay.disconnectAgent(tunnelId, 4000, 'heartbeat timeout');
         this.states.delete(tunnelId);
       }
     }

package/src/server/relay.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, test } from 'bun:test';
+import { TunnelRelay } from './relay';
+
+function fakeWs(closes: Array<{ code?: number; reason?: string }> = []): WebSocket {
+  return {
+    readyState: WebSocket.OPEN,
+    send: () => {},
+    close: (code?: number, reason?: string) => closes.push({ code, reason }),
+  } as unknown as WebSocket;
+}
+
+describe('TunnelRelay connection lifecycle', () => {
+  test('updates live metadata without exposing session signing state', () => {
+    const relay = new TunnelRelay();
+    relay.registerAgent('tnl_1', fakeWs(), 'signing-key', {
+      accountId: 'acct_1',
+      capabilities: ['filesystem'],
+    });
+
+    expect(relay.updateAgentMetadata('tnl_1', { capabilities: ['desktop'] })).toBe(true);
+    expect(relay.getAgentMetadata('tnl_1')).toEqual({
+      accountId: 'acct_1',
+      capabilities: ['desktop'],
+    });
+  });
+
+  test('stale close from a replaced socket does not unregister the active agent', () => {
+    const relay = new TunnelRelay();
+    const closes: Array<{ code?: number; reason?: string }> = [];
+    const first = fakeWs(closes);
+    const second = fakeWs();
+
+    relay.registerAgent('tnl_1', first, 'signing-key-1', {
+      accountId: 'acct_1',
+    });
+    relay.registerAgent('tnl_1', second, 'signing-key-2', {
+      accountId: 'acct_1',
+    });
+
+    const removed = relay.unregisterAgent('tnl_1', first);
+
+    expect(removed).toBe(false);
+    expect(relay.isConnected('tnl_1')).toBe(true);
+    expect(relay.getConnectedCount()).toBe(1);
+    expect(closes).toEqual([{ code: 4004, reason: 'replaced by another agent process' }]);
+  });
+
+  test('close from the active socket unregisters the agent and emits metadata', () => {
+    const relay = new TunnelRelay();
+    const ws = fakeWs();
+    const events: unknown[] = [];
+    relay.on('agent:disconnect', (event) => events.push(event));
+
+    relay.registerAgent('tnl_1', ws, 'signing-key', { accountId: 'acct_1' });
+    const removed = relay.unregisterAgent('tnl_1', ws);
+
+    expect(removed).toBe(true);
+    expect(relay.isConnected('tnl_1')).toBe(false);
+    expect(events).toEqual([{ tunnelId: 'tnl_1', metadata: { accountId: 'acct_1' } }]);
+  });
+
+  test('rejects an outgoing RPC larger than the configured byte limit', async () => {
+    const relay = new TunnelRelay({ maxWsMessageSize: 128 });
+    relay.registerAgent('tnl_1', fakeWs(), 'signing-key', {
+      accountId: 'acct_1',
+    });
+
+    await expect(relay.relayRPC('tnl_1', 'fs.write', { content: 'é'.repeat(128) })).rejects.toThrow(
+      'exceeds the maximum tunnel message size',
+    );
+  });
+});

package/src/server/relay.ts

@@ -2,7 +2,6 @@ import { EventEmitter } from 'events';
 import { signMessage, verifyMessageSignature } from '../shared/crypto';
 import {
   type JsonRpcRequest,
-  type JsonRpcResponse,
   type JsonRpcNotification,
   type PendingRPC,
   type RelayRpcOptions,
@@ -60,7 +59,9 @@ export class TunnelRelay extends EventEmitter {
           this.pendingRPCs.delete(requestId);
         }
       }
-      try { existing.ws.close(1000, 'replaced by new connection'); } catch {}
+      // A normal close makes the displaced agent reconnect and replace the new
+      // socket forever. Use a terminal code so one process wins deterministically.
+      try { existing.ws.close(4004, 'replaced by another agent process'); } catch {}
       this.emitEvent('connection:replaced', { tunnelId });
     }
 
@@ -69,7 +70,16 @@ export class TunnelRelay extends EventEmitter {
     console.log(`[tunnel-relay] Agent registered: ${tunnelId} (total: ${this.agents.size})`);
   }
 
-  unregisterAgent(tunnelId: string): void {
+  unregisterAgent(tunnelId: string, ws?: WebSocket): boolean {
+    const existing = this.agents.get(tunnelId);
+    if (!existing) {
+      return false;
+    }
+    if (ws && existing && existing.ws !== ws) {
+      return false;
+    }
+
+    const metadata = existing?.metadata;
     this.agents.delete(tunnelId);
 
     for (const [requestId, pending] of this.pendingRPCs) {
@@ -83,8 +93,19 @@ export class TunnelRelay extends EventEmitter {
       }
     }
 
-    this.emitEvent('agent:disconnect', { tunnelId });
+    this.emitEvent('agent:disconnect', { tunnelId, metadata });
     console.log(`[tunnel-relay] Agent unregistered: ${tunnelId} (total: ${this.agents.size})`);
+    return true;
+  }
+
+  disconnectAgent(tunnelId: string, code = 1000, reason = 'disconnected by server'): boolean {
+    const agent = this.agents.get(tunnelId);
+    if (!agent) return false;
+    const removed = this.unregisterAgent(tunnelId, agent.ws);
+    if (removed) {
+      try { agent.ws.close(code, reason); } catch {}
+    }
+    return removed;
   }
 
   isConnected(tunnelId: string): boolean {
@@ -111,7 +132,14 @@ export class TunnelRelay extends EventEmitter {
     return this.agents.get(tunnelId)?.metadata;
   }
 
-  handleAgentMessage(tunnelId: string, raw: string | Buffer): void {
+  updateAgentMetadata(tunnelId: string, patch: Record<string, unknown>): boolean {
+    const agent = this.agents.get(tunnelId);
+    if (!agent) return false;
+    agent.metadata = { ...(agent.metadata ?? {}), ...patch };
+    return true;
+  }
+
+  handleAgentMessage(tunnelId: string, ws: WebSocket, raw: string | Buffer): void {
     let msg: any;
     try {
       msg = JSON.parse(typeof raw === 'string' ? raw : raw.toString('utf-8'));
@@ -122,6 +150,10 @@ export class TunnelRelay extends EventEmitter {
 
     // Verify HMAC signature on ALL messages from agent (including pong)
     const agent = this.agents.get(tunnelId);
+    if (!agent || agent.ws !== ws) {
+      console.warn(`[tunnel-relay] Message from inactive socket for ${tunnelId}, discarding`);
+      return;
+    }
     if (agent && msg._sig !== undefined && msg._nonce !== undefined) {
       if (msg._nonce <= agent.lastResponseNonce) {
         console.warn(`[tunnel-relay] Replay detected from agent ${tunnelId}: nonce ${msg._nonce} <= ${agent.lastResponseNonce}`);
@@ -141,7 +173,7 @@ export class TunnelRelay extends EventEmitter {
       console.warn(`[tunnel-relay] Unsigned message from agent ${tunnelId}, discarding`);
       return;
     }
-   
+
     if ('method' in msg && msg.method === 'tunnel.pong') {
       this.emitEvent('message:pong', { tunnelId, params: msg.params });
       return;
@@ -153,7 +185,7 @@ export class TunnelRelay extends EventEmitter {
     }
 
     const pending = this.pendingRPCs.get(msg.id);
-    if (!pending) {
+    if (!pending || pending.tunnelId !== tunnelId) {
       return;
     }
 
@@ -220,6 +252,13 @@ export class TunnelRelay extends EventEmitter {
     const payload = JSON.stringify(request);
     const sig = signMessage(agent.signingKey, payload, nonce);
     const signedRequest = { ...request, _sig: sig, _nonce: nonce };
+    const encodedRequest = JSON.stringify(signedRequest);
+    if (Buffer.byteLength(encodedRequest, 'utf8') > this.config.maxWsMessageSize) {
+      throw new TunnelRelayError(
+        TunnelErrorCode.LOCAL_ERROR,
+        `RPC request exceeds the maximum tunnel message size for ${method}`,
+      );
+    }
 
     this.emitEvent('rpc:request', { tunnelId, method, requestId });
 
@@ -244,7 +283,7 @@ export class TunnelRelay extends EventEmitter {
       });
 
       try {
-        agent.ws.send(JSON.stringify(signedRequest));
+        agent.ws.send(encodedRequest);
       } catch (err) {
         clearTimeout(timer);
         this.pendingRPCs.delete(requestId);
@@ -272,7 +311,9 @@ export class TunnelRelay extends EventEmitter {
     const signedNotification = { ...notification, _sig: sig, _nonce: nonce };
 
     try {
-      agent.ws.send(JSON.stringify(signedNotification));
+      const encoded = JSON.stringify(signedNotification);
+      if (Buffer.byteLength(encoded, 'utf8') > this.config.maxWsMessageSize) return false;
+      agent.ws.send(encoded);
       return true;
     } catch {
       return false;

package/src/server/server.test.ts

@@ -0,0 +1,33 @@
+import { afterEach, describe, expect, test } from 'bun:test';
+
+import { startTunnelServer, type TunnelServer } from './server';
+
+let server: TunnelServer | null = null;
+
+afterEach(() => {
+  server?.stop();
+  server = null;
+});
+
+describe('standalone tunnel WebSocket boundary', () => {
+  test('rejects browser-origin WebSocket upgrades before authentication', async () => {
+    server = startTunnelServer({ port: 0 });
+    const websocketKey = Buffer.from('fixed-test-nonce').toString('base64');
+    const response = await fetch(
+      `http://127.0.0.1:${server.port}/ws?tunnelId=00000000-0000-4000-8000-000000000001`,
+      {
+        headers: {
+          connection: 'Upgrade',
+          upgrade: 'websocket',
+          origin: 'https://attacker.example',
+          'sec-websocket-key': websocketKey,
+          'sec-websocket-version': '13',
+        },
+      },
+    );
+
+    expect(response.status).toBe(403);
+    expect(await response.json()).toEqual({ error: 'Browser tunnel WebSockets are not allowed' });
+    expect(server.relay.getConnectedCount()).toBe(0);
+  });
+});

package/src/server/server.ts

@@ -18,6 +18,7 @@ import { createWsHandlers } from './ws-handler';
 import type { TunnelServerConfig } from '../shared/types';
 
 export interface TunnelServer {
+  port: number;
   app: Hono;
   relay: TunnelRelay;
   heartbeat: HeartbeatManager;
@@ -70,13 +71,25 @@ export function startTunnelServer(config?: TunnelServerConfig): TunnelServer {
       if (url.pathname === '/ws') {
         const tunnelId = url.searchParams.get('tunnelId');
 
-        if (!tunnelId) {
-          return new Response(JSON.stringify({ error: 'Missing tunnelId' }), {
+        if (!tunnelId || !/^[A-Za-z0-9._:-]{1,128}$/.test(tunnelId)) {
+          return new Response(JSON.stringify({ error: 'Invalid tunnelId' }), {
             status: 400,
             headers: { 'Content-Type': 'application/json' },
           });
         }
 
+        if (req.headers.has('origin')) {
+          return new Response(
+            JSON.stringify({
+              error: 'Browser tunnel WebSockets are not allowed',
+            }),
+            {
+              status: 403,
+              headers: { 'Content-Type': 'application/json' },
+            },
+          );
+        }
+
         const success = server.upgrade(req, {
           data: { tunnelId } as any,
         });
@@ -94,15 +107,15 @@ export function startTunnelServer(config?: TunnelServerConfig): TunnelServer {
         wsHandlers.onOpen(ws.data.tunnelId, ws);
       },
       message(ws: any, message: string | Buffer) {
-        wsHandlers.onMessage(ws.data.tunnelId, message);
+        wsHandlers.onMessage(ws.data.tunnelId, ws, message);
       },
       close(ws: any) {
-        wsHandlers.onClose(ws.data.tunnelId);
+        wsHandlers.onClose(ws.data.tunnelId, ws);
       },
     },
   });
 
-  console.log(`[agent-tunnel] Server listening on port ${port}`);
+  console.log(`[agent-tunnel] Server listening on port ${bunServer.port}`);
 
   const stop = () => {
     heartbeat.stop();
@@ -110,5 +123,12 @@ export function startTunnelServer(config?: TunnelServerConfig): TunnelServer {
     bunServer.stop();
   };
 
-  return { app, relay, heartbeat, wsHandlers, stop };
+  return {
+    port: bunServer.port ?? port,
+    app,
+    relay,
+    heartbeat,
+    wsHandlers,
+    stop,
+  };
 }