Add streaming relay protocol — forward chunks as they arrive

electron-agent/agent/TunnelAgent.ts

@@ -18,13 +18,6 @@ interface RelayRequest
   body:    string; // base64
 }
 
-interface RelayResponse
-{
-  reqId:   string;
-  status:  number;
-  headers: Record<string, string>;
-  body:    string; // base64
-}
 
 export class TunnelAgent
 {
@@ -84,6 +77,8 @@ export class TunnelAgent
     );
   }
 
+  private _send( msg: object ): void { this.ws?.send( JSON.stringify( msg ) ); }
+
   private _forward( data: Buffer ): void
   {
     let req: RelayRequest;
@@ -100,40 +95,32 @@ export class TunnelAgent
       headers:  { ...req.headers, 'content-length': String( bodyBuf.length ) },
     };
 
-    const chunks: Buffer[] = [];
-
     const localReq = http.request( opts, localRes =>
     {
-      localRes.on( 'data', ( chunk: Buffer ) => chunks.push( Buffer.from( chunk ) ) );
-      localRes.on( 'end', () =>
+      const respHeaders: Record<string, string> = {};
+      for ( const [ k, v ] of Object.entries( localRes.headers ) )
       {
-        const respHeaders: Record<string, string> = {};
-        for ( const [ k, v ] of Object.entries( localRes.headers ) )
-        {
-          if ( typeof v === 'string' )          respHeaders[ k ] = v;
-          else if ( Array.isArray( v ) )        respHeaders[ k ] = v.join( ', ' );
-        }
+        if ( typeof v === 'string' )   respHeaders[ k ] = v;
+        else if ( Array.isArray( v ) ) respHeaders[ k ] = v.join( ', ' );
+      }
 
-        const resp: RelayResponse = {
-          reqId:   req.reqId,
-          status:  localRes.statusCode ?? 200,
-          headers: respHeaders,
-          body:    Buffer.concat( chunks ).toString( 'base64' ),
-        };
-        this.ws?.send( JSON.stringify( resp ) );
-      } );
+      // Send headers immediately — starts streaming
+      this._send( { type: 'res_start', reqId: req.reqId, status: localRes.statusCode ?? 200, headers: respHeaders } );
+
+      localRes.on( 'data', ( chunk: Buffer ) =>
+        this._send( { type: 'res_data', reqId: req.reqId, chunk: Buffer.from( chunk ).toString( 'base64' ) } )
+      );
+
+      localRes.on( 'end', () =>
+        this._send( { type: 'res_end', reqId: req.reqId } )
+      );
     } );
 
     localReq.on( 'error', err =>
     {
-      const errResp: RelayResponse = {
-        reqId:   req.reqId,
-        status:  502,
-        headers: { 'content-type': 'application/json' },
-        body:    Buffer.from( JSON.stringify( { error: 'Local service error', detail: err.message } ) )
-                       .toString( 'base64' ),
-      };
-      this.ws?.send( JSON.stringify( errResp ) );
+      this._send( { type: 'res_start', reqId: req.reqId, status: 502, headers: { 'content-type': 'application/json' } } );
+      this._send( { type: 'res_data',  reqId: req.reqId, chunk: Buffer.from( JSON.stringify( { error: err.message } ) ).toString( 'base64' ) } );
+      this._send( { type: 'res_end',   reqId: req.reqId } );
     } );
 
     if ( bodyBuf.length > 0 ) localReq.write( bodyBuf );

source/server/relay/pending.ts

@@ -1,27 +1,63 @@
 export interface RelayResponse
 {
-  reqId: string;
-  status: number;
+  reqId:   string;
+  status:  number;
   headers: Record<string, string>;
-  body: string; // base64-encoded
+  body:    string; // base64 — legacy single-shot
 }
 
-const callbacks = new Map<string, ( r: RelayResponse ) => void>();
+export type AgentMsg =
+  | { type: 'res_start'; reqId: string; status: number; headers: Record<string, string> }
+  | { type: 'res_data';  reqId: string; chunk: string } // base64
+  | { type: 'res_end';   reqId: string }
+  | RelayResponse; // legacy: no type field
 
-export function addPending( reqId: string, cb: ( r: RelayResponse ) => void ): void
+export interface StreamHandler
 {
-  callbacks.set( reqId, cb );
+  onStart: ( status: number, headers: Record<string, string> ) => void;
+  onData:  ( chunk: Buffer ) => void;
+  onEnd:   () => void;
 }
 
-export function resolvePending( resp: RelayResponse ): void
+const pending = new Map<string, StreamHandler>();
+
+export function addPending( reqId: string, handler: StreamHandler ): void
 {
-  const cb = callbacks.get( resp.reqId );
-  if ( !cb ) return;
-  callbacks.delete( resp.reqId );
-  cb( resp );
+  pending.set( reqId, handler );
 }
 
 export function removePending( reqId: string ): void
 {
-  callbacks.delete( reqId );
+  pending.delete( reqId );
+}
+
+export function dispatchAgentMsg( msg: AgentMsg ): void
+{
+  const handler = pending.get( msg.reqId );
+  if ( !handler ) return;
+
+  if ( 'type' in msg )
+  {
+    if ( msg.type === 'res_start' )
+    {
+      handler.onStart( msg.status, msg.headers );
+    }
+    else if ( msg.type === 'res_data' )
+    {
+      handler.onData( Buffer.from( msg.chunk, 'base64' ) );
+    }
+    else if ( msg.type === 'res_end' )
+    {
+      pending.delete( msg.reqId );
+      handler.onEnd();
+    }
+  }
+  else
+  {
+    // Legacy single-shot RelayResponse — forward as stream
+    handler.onStart( msg.status, msg.headers );
+    handler.onData( Buffer.from( msg.body, 'base64' ) );
+    pending.delete( msg.reqId );
+    handler.onEnd();
+  }
 }

source/server/routes/agent.ts

@@ -4,7 +4,7 @@ import jwt from 'jsonwebtoken';
 import { AuthPayload } from '../middleware/requireAuth';
 import { getTunnel } from '../db';
 import { registry } from '../relay/TunnelRegistry';
-import { resolvePending } from '../relay/pending';
+import { dispatchAgentMsg } from '../relay/pending';
 
 export function handleAgentUpgrade( req: IncomingMessage, ws: WebSocket ): void
 {
@@ -36,7 +36,7 @@ export function handleAgentUpgrade( req: IncomingMessage, ws: WebSocket ): void
 
   ws.on( 'message', ( data ) =>
   {
-    try { resolvePending( JSON.parse( data.toString() ) ); }
+    try { dispatchAgentMsg( JSON.parse( data.toString() ) ); }
     catch { /* ignore malformed messages */ }
   } );
 

source/server/routes/proxy.ts

@@ -4,15 +4,15 @@ import crypto from 'crypto';
 import { getTunnel, TunnelConfig } from '../db';
 import { registry } from '../relay/TunnelRegistry';
 import { AuthPayload, extractBearer } from '../middleware/requireAuth';
-import { addPending, removePending, RelayResponse } from '../relay/pending';
+import { addPending, removePending } from '../relay/pending';
 
 export interface RelayRequest
 {
-  reqId: string;
-  method: string;
-  path: string;
+  reqId:   string;
+  method:  string;
+  path:    string;
   headers: Record<string, string>;
-  body: string; // base64-encoded
+  body:    string; // base64-encoded
 }
 
 const router = Router();
@@ -48,10 +48,10 @@ router.all( [ '/:tunnelId', '/:tunnelId/*' ], async ( req, res ) =>
   if ( !ws ) { res.status( 503 ).json( { error: 'Tunnel agent not connected' } ); return; }
 
   // Strip the /:tunnelId prefix to get the path to forward
-  const prefix = '/' + req.params.tunnelId;
+  const prefix      = '/' + req.params.tunnelId;
   const forwardPath = req.url.slice( prefix.length ) || '/';
 
-  // Buffer the raw request body (no JSON middleware on /t routes)
+  // Buffer the raw request body
   const chunks: Buffer[] = [];
   await new Promise<void>( ( resolve, reject ) =>
   {
@@ -74,39 +74,41 @@ router.all( [ '/:tunnelId', '/:tunnelId/*' ], async ( req, res ) =>
 
   const relayReq: RelayRequest = {
     reqId,
-    method: req.method,
-    path: forwardPath,
+    method:  req.method,
+    path:    forwardPath,
     headers,
-    body: bodyBuf.toString( 'base64' ),
+    body:    bodyBuf.toString( 'base64' ),
   };
 
-  const TIMEOUT_MS = 30_000;
-  const response = await new Promise<RelayResponse | null>( ( resolve ) =>
+  // 2-minute timeout for the first response chunk — reset once streaming starts
+  await new Promise<void>( ( resolve ) =>
   {
     const timer = setTimeout( () =>
     {
       removePending( reqId );
-      resolve( null );
-    }, TIMEOUT_MS );
+      if ( !res.headersSent ) res.status( 504 ).json( { error: 'Agent timed out' } );
+      resolve();
+    }, 120_000 );
 
-    addPending( reqId, ( resp ) =>
+    addPending( reqId,
     {
-      clearTimeout( timer );
-      resolve( resp );
+      onStart: ( status, respHeaders ) =>
+      {
+        clearTimeout( timer ); // streaming started — no more timeout
+        res.status( status );
+        for ( const [ key, val ] of Object.entries( respHeaders ) )
+        {
+          if ( key.toLowerCase() === 'transfer-encoding' ) continue;
+          res.setHeader( key, val );
+        }
+        res.flushHeaders();
+      },
+      onData: ( chunk ) => res.write( chunk ),
+      onEnd:  () => { res.end(); resolve(); },
     } );
 
     ws.send( JSON.stringify( relayReq ) );
   } );
-
-  if ( !response ) { res.status( 504 ).json( { error: 'Agent timed out' } ); return; }
-
-  res.status( response.status );
-  for ( const [ key, val ] of Object.entries( response.headers ) )
-  {
-    if ( key.toLowerCase() === 'transfer-encoding' ) continue;
-    res.setHeader( key, val );
-  }
-  res.end( Buffer.from( response.body, 'base64' ) );
 } );
 
 export default router;