@agentproto/acp 0.1.0 → 0.3.0

package/dist/tunnel/index.mjs

@@ -58,6 +58,8 @@ function encodeData(bytes) {
 function decodeData(data) {
   return Buffer.from(data, "base64");
 }
+var DEFAULT_WS_DIAL_TIMEOUT_MS = 1e4;
+var DEFAULT_HTTP_FORWARD_TIMEOUT_MS = 3e4;
 function createTunnelServer(opts) {
   const children = /* @__PURE__ */ new Map();
   const ptyProcs = /* @__PURE__ */ new Map();
@@ -68,7 +70,8 @@ function createTunnelServer(opts) {
     version: TUNNEL_VERSION,
     capabilities: {
       pty: opts.pty === true,
-      wsForward: opts.dialUpstreamWs !== void 0 && !!opts.httpUpstream
+      wsForward: opts.dialUpstreamWs !== void 0 && !!opts.httpUpstream,
+      ...opts.tools && opts.tools.length ? { tools: opts.tools } : {}
     },
     label: opts.label,
     daemon: {
@@ -383,22 +386,69 @@ function createTunnelServer(opts) {
       if (HOP_BY_HOP.has(k.toLowerCase())) continue;
       headers[k] = v;
     }
-    const base = opts.httpUpstream.replace(/^http(s?):\/\//, "ws$1://");
+    let httpBase = opts.httpUpstream;
+    if (req.upstream) {
+      const resolved = await opts.resolveWsUpstream?.(req.upstream);
+      if (!resolved) {
+        reject(
+          502,
+          "unknown_ws_upstream",
+          `No WS upstream registered for '${req.upstream}'.`
+        );
+        return;
+      }
+      httpBase = resolved;
+    }
+    const base = httpBase.replace(/^http(s?):\/\//, "ws$1://");
     const url = `${base.replace(/\/$/, "")}${req.path}`;
+    const dialTimeoutMs = opts.wsDialTimeoutMs ?? DEFAULT_WS_DIAL_TIMEOUT_MS;
+    const dialAbort = new AbortController();
+    let dialTimedOut = false;
+    let dialTimer = null;
+    const dialPromise = opts.dialUpstreamWs({
+      url,
+      ...req.protocols && req.protocols.length > 0 ? { protocols: req.protocols } : {},
+      headers,
+      signal: dialAbort.signal
+    });
+    const timeoutPromise = new Promise((_resolve, rej) => {
+      dialTimer = setTimeout(() => {
+        dialTimedOut = true;
+        dialAbort.abort(new Error("ws_dial_timeout"));
+        rej(new Error(`upstream WS dial timed out after ${dialTimeoutMs}ms`));
+      }, dialTimeoutMs);
+      if (typeof dialTimer !== "number") dialTimer.unref();
+    });
     let upstream;
     try {
-      upstream = await opts.dialUpstreamWs({
-        url,
-        ...req.protocols && req.protocols.length > 0 ? { protocols: req.protocols } : {},
-        headers
-      });
+      upstream = await Promise.race([dialPromise, timeoutPromise]);
     } catch (err) {
+      if (dialTimer) clearTimeout(dialTimer);
+      if (dialTimedOut) {
+        dialPromise.then(
+          (sock) => {
+            try {
+              sock.close(1001, "dial_timeout");
+            } catch {
+            }
+          },
+          () => {
+          }
+        );
+        reject(
+          504,
+          "upstream_ws_timeout",
+          `Upstream did not complete the WS upgrade within ${dialTimeoutMs}ms.`
+        );
+        return;
+      }
       const message = err instanceof Error ? err.message : String(err);
       const m = message.match(/Unexpected server response: (\d+)/);
       const status = m ? Number.parseInt(m[1], 10) : 502;
       reject(status, "upstream_ws_failed", message);
       return;
     }
+    if (dialTimer) clearTimeout(dialTimer);
     if (closed) {
       try {
         upstream.close(1001, "tunnel_closed");
@@ -501,7 +551,7 @@ function createTunnelServer(opts) {
     }
     const url = `${opts.httpUpstream.replace(/\/$/, "")}${req.path}`;
     const controller = new AbortController();
-    const timeoutMs = req.timeoutMs ?? 3e4;
+    const timeoutMs = req.timeoutMs ?? opts.httpForwardTimeoutMs ?? DEFAULT_HTTP_FORWARD_TIMEOUT_MS;
     let timer = setTimeout(
       () => controller.abort(),
       timeoutMs
@@ -537,9 +587,41 @@ function createTunnelServer(opts) {
           headers: outHeaders
         });
         const reader = upstreamRes.body.getReader();
+        const idleMs = opts.httpStreamIdleTimeoutMs;
+        const IDLE = /* @__PURE__ */ Symbol("idle");
+        const readChunk = async () => {
+          if (!idleMs) return reader.read();
+          let idleTimer = null;
+          const idle = new Promise((resolve) => {
+            idleTimer = setTimeout(() => resolve(IDLE), idleMs);
+            if (typeof idleTimer !== "number") idleTimer.unref();
+          });
+          try {
+            return await Promise.race([reader.read(), idle]);
+          } finally {
+            if (idleTimer) clearTimeout(idleTimer);
+          }
+        };
         try {
           while (true) {
-            const { value, done } = await reader.read();
+            const chunk = await readChunk();
+            if (chunk === IDLE) {
+              try {
+                await reader.cancel();
+              } catch {
+              }
+              opts.sink.send({
+                t: "http_response_chunk",
+                reqId: req.reqId,
+                end: true,
+                error: {
+                  code: "upstream_stream_idle_timeout",
+                  message: `No data from upstream for ${idleMs}ms.`
+                }
+              });
+              return;
+            }
+            const { value, done } = chunk;
             if (done) {
               opts.sink.send({
                 t: "http_response_chunk",
@@ -1031,6 +1113,7 @@ function createTunnelClient(opts) {
         t: "ws_open",
         reqId,
         path: req.path,
+        ...req.upstream ? { upstream: req.upstream } : {},
         ...req.headers ? { headers: req.headers } : {},
         ...req.protocols && req.protocols.length > 0 ? { protocols: req.protocols } : {}
       };
@@ -1280,6 +1363,6 @@ function wrapWebSocket(ws) {
   };
 }
 
-export { TUNNEL_VERSION, createTunnelClient, createTunnelServer, decodeData, encodeData, encodeFrame, parseFrame, wrapWebSocket };
+export { DEFAULT_HTTP_FORWARD_TIMEOUT_MS, DEFAULT_WS_DIAL_TIMEOUT_MS, TUNNEL_VERSION, createTunnelClient, createTunnelServer, decodeData, encodeData, encodeFrame, parseFrame, wrapWebSocket };
 //# sourceMappingURL=index.mjs.map
 //# sourceMappingURL=index.mjs.map