@kortix/agent-tunnel 0.15.0 → 0.15.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@kortix/agent-tunnel",
-  "version": "0.15.0",
+  "version": "0.15.2",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/kortix-ai/CompanyOS.git",

package/src/server/relay.ts

@@ -11,8 +11,15 @@ import {
   TunnelErrorCode,
 } from '../shared/types';
 
+/** What the relay uses of an agent's socket. A DOM `WebSocket` and Bun's `ServerWebSocket` both fit it. */
+export interface AgentSocket {
+  readonly readyState: number;
+  send(data: string): void;
+  close(code?: number, reason?: string): void;
+}
+
 interface AgentConnection {
-  ws: WebSocket;
+  ws: AgentSocket;
   signingKey: string;
   nonce: number;
   lastResponseNonce: number;
@@ -43,7 +50,7 @@ export class TunnelRelay extends EventEmitter {
 
   registerAgent(
     tunnelId: string,
-    ws: WebSocket,
+    ws: AgentSocket,
     signingKey: string,
     metadata?: Record<string, unknown>,
   ): void {
@@ -70,7 +77,7 @@ export class TunnelRelay extends EventEmitter {
     console.log(`[tunnel-relay] Agent registered: ${tunnelId} (total: ${this.agents.size})`);
   }
 
-  unregisterAgent(tunnelId: string, ws?: WebSocket): boolean {
+  unregisterAgent(tunnelId: string, ws?: AgentSocket): boolean {
     const existing = this.agents.get(tunnelId);
     if (!existing) {
       return false;
@@ -139,7 +146,7 @@ export class TunnelRelay extends EventEmitter {
     return true;
   }
 
-  handleAgentMessage(tunnelId: string, ws: WebSocket, raw: string | Buffer): void {
+  handleAgentMessage(tunnelId: string, ws: AgentSocket, raw: string | Buffer): void {
     let msg: any;
     try {
       msg = JSON.parse(typeof raw === 'string' ? raw : raw.toString('utf-8'));

package/src/server/server.ts

@@ -63,7 +63,7 @@ export function startTunnelServer(config?: TunnelServerConfig): TunnelServer {
 
   heartbeat.start();
 
-  const bunServer = Bun.serve({
+  const bunServer = Bun.serve<{ tunnelId: string }>({
     port,
     fetch(req, server) {
       const url = new URL(req.url);
@@ -91,7 +91,7 @@ export function startTunnelServer(config?: TunnelServerConfig): TunnelServer {
         }
 
         const success = server.upgrade(req, {
-          data: { tunnelId } as any,
+          data: { tunnelId },
         });
         if (success) return undefined;
 
@@ -103,13 +103,13 @@ export function startTunnelServer(config?: TunnelServerConfig): TunnelServer {
 
     websocket: {
       idleTimeout: 0,
-      open(ws: any) {
+      open(ws) {
         wsHandlers.onOpen(ws.data.tunnelId, ws);
       },
-      message(ws: any, message: string | Buffer) {
+      message(ws, message: string | Buffer) {
         wsHandlers.onMessage(ws.data.tunnelId, ws, message);
       },
-      close(ws: any) {
+      close(ws) {
         wsHandlers.onClose(ws.data.tunnelId, ws);
       },
     },

package/src/server/ws-handler.ts

@@ -1,4 +1,4 @@
-import type { TunnelRelay } from './relay';
+import type { AgentSocket, TunnelRelay } from './relay';
 import type { HeartbeatManager } from './heartbeat';
 import type { TunnelAuthMessage, AuthResult } from '../shared/types';
 
@@ -14,9 +14,9 @@ export interface WsHandlerOptions {
 }
 
 export interface WsHandlers {
-  onOpen(tunnelId: string, ws: WebSocket): void;
-  onMessage(tunnelId: string, ws: WebSocket, message: string | Buffer): void;
-  onClose(tunnelId: string, ws?: WebSocket): void;
+  onOpen(tunnelId: string, ws: AgentSocket): void;
+  onMessage(tunnelId: string, ws: AgentSocket, message: string | Buffer): void;
+  onClose(tunnelId: string, ws?: AgentSocket): void;
 }
 
 export function createWsHandlers(relay: TunnelRelay, opts?: WsHandlerOptions): WsHandlers {
@@ -26,7 +26,7 @@ export function createWsHandlers(relay: TunnelRelay, opts?: WsHandlerOptions): W
   const authTimeoutMs = opts?.authTimeoutMs ?? 10_000;
 
   const pendingConnections = new Map<
-    WebSocket,
+    AgentSocket,
     {
       tunnelId: string;
       timer: ReturnType<typeof setTimeout>;
@@ -35,7 +35,7 @@ export function createWsHandlers(relay: TunnelRelay, opts?: WsHandlerOptions): W
   >();
 
   return {
-    onOpen(tunnelId: string, ws: WebSocket) {
+    onOpen(tunnelId: string, ws: AgentSocket) {
       const timer = setTimeout(() => {
         pendingConnections.delete(ws);
         try {
@@ -46,7 +46,7 @@ export function createWsHandlers(relay: TunnelRelay, opts?: WsHandlerOptions): W
       pendingConnections.set(ws, { tunnelId, timer, authenticating: false });
     },
 
-    async onMessage(tunnelId: string, ws: WebSocket, message: string | Buffer) {
+    async onMessage(tunnelId: string, ws: AgentSocket, message: string | Buffer) {
       const msgStr = typeof message === 'string' ? message : message.toString('utf-8');
       const msgSize =
         typeof message === 'string'
@@ -159,7 +159,7 @@ export function createWsHandlers(relay: TunnelRelay, opts?: WsHandlerOptions): W
       relay.handleAgentMessage(tunnelId, ws, message);
     },
 
-    onClose(tunnelId: string, ws?: WebSocket) {
+    onClose(tunnelId: string, ws?: AgentSocket) {
       const pending = ws ? pendingConnections.get(ws) : undefined;
       if (pending) {
         clearTimeout(pending.timer);