Handle upstream stream drops

src/upstream.ts

@@ -7,11 +7,15 @@
  * header without coupling to a transport class.
  */
 const JSON_RPC = "2.0";
+const DEFAULT_RPC_TIMEOUT_MS = 30_000;
+const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 120_000;
 
 export interface UpstreamConfig {
   name: string;
   url: string;
   token: string | null;
+  requestTimeoutMs?: number;
+  streamIdleTimeoutMs?: number;
   /** Forwarded to the upstream as ``params._meta.cp_jwt`` on tools/call.
    *  Lets the upstream agent act on the caller's behalf (file/agent CRUD
    *  on /v1/me/*) — same path the platform orchestrator uses. */
@@ -63,32 +67,53 @@ export class UpstreamAgent {
 
   private async rpc<T>(method: string, params: unknown): Promise<T> {
     const id = _nextId++;
+    const timeoutMs = this.cfg.requestTimeoutMs ?? DEFAULT_RPC_TIMEOUT_MS;
+    const controller = new AbortController();
+    let timedOut = false;
+    const timeout = setTimeout(() => {
+      timedOut = true;
+      controller.abort();
+    }, timeoutMs);
     const headers: Record<string, string> = {
       "content-type": "application/json",
       accept: "application/json",
     };
     if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`;
 
-    const resp = await fetch(`${this.cfg.url}/mcp`, {
-      method: "POST",
-      headers,
-      body: JSON.stringify({ jsonrpc: JSON_RPC, id, method, params }),
-    });
+    try {
+      const resp = await fetch(`${this.cfg.url}/mcp`, {
+        method: "POST",
+        headers,
+        body: JSON.stringify({ jsonrpc: JSON_RPC, id, method, params }),
+        signal: controller.signal,
+      });
 
-    if (!resp.ok) {
-      const text = await resp.text().catch(() => "");
-      throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText);
-    }
+      if (!resp.ok) {
+        const text = await resp.text().catch(() => "");
+        throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText);
+      }
 
-    const env: any = await resp.json();
-    if (env?.error) {
-      throw new UpstreamError(
-        this.cfg.name,
-        500,
-        `${env.error.code}: ${env.error.message}`,
-      );
+      const env: any = await resp.json();
+      if (env?.error) {
+        throw new UpstreamError(
+          this.cfg.name,
+          500,
+          `${env.error.code}: ${env.error.message}`,
+        );
+      }
+      return env?.result as T;
+    } catch (err) {
+      if (timedOut) {
+        throw new UpstreamError(
+          this.cfg.name,
+          504,
+          `request timed out after ${timeoutMs}ms`,
+        );
+      }
+      throw err;
+    } finally {
+      clearTimeout(timeout);
     }
-    return env?.result as T;
   }
 
   async listTools(): Promise<UpstreamTool[]> {
@@ -136,22 +161,51 @@ export class UpstreamAgent {
     }) => void | Promise<void>,
   ): Promise<UpstreamCallResult> {
     const id = _nextId++;
+    const idleTimeoutMs =
+      this.cfg.streamIdleTimeoutMs ?? DEFAULT_STREAM_IDLE_TIMEOUT_MS;
+    const controller = new AbortController();
+    let timedOut = false;
+    let idleTimer: ReturnType<typeof setTimeout> | undefined;
+    const resetIdleTimer = () => {
+      if (idleTimer) clearTimeout(idleTimer);
+      idleTimer = setTimeout(() => {
+        timedOut = true;
+        controller.abort();
+      }, idleTimeoutMs);
+    };
+    resetIdleTimer();
     const headers: Record<string, string> = {
       "content-type": "application/json",
       accept: "text/event-stream, application/json",
     };
     if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`;
 
-    const resp = await fetch(`${this.cfg.url}/mcp`, {
-      method: "POST",
-      headers,
-      body: JSON.stringify({
-        jsonrpc: JSON_RPC, id, method: "tools/call", params,
-      }),
-    });
+    let resp: Response;
+    try {
+      resp = await fetch(`${this.cfg.url}/mcp`, {
+        method: "POST",
+        headers,
+        body: JSON.stringify({
+          jsonrpc: JSON_RPC, id, method: "tools/call", params,
+        }),
+        signal: controller.signal,
+      });
+    } catch (err) {
+      clearTimeout(idleTimer);
+      if (timedOut) {
+        throw new UpstreamError(
+          this.cfg.name,
+          504,
+          `stream idle timeout after ${idleTimeoutMs}ms`,
+        );
+      }
+      throw err;
+    }
+    resetIdleTimer();
 
     if (!resp.ok) {
       const text = await resp.text().catch(() => "");
+      clearTimeout(idleTimer);
       throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText);
     }
 
@@ -160,15 +214,18 @@ export class UpstreamAgent {
       // Server fell back to JSON (older a2a-pack or non-SSE upstream).
       const env: any = await resp.json();
       if (env?.error) {
+        clearTimeout(idleTimer);
         throw new UpstreamError(
           this.cfg.name, 500,
           `${env.error.code}: ${env.error.message}`,
         );
       }
+      clearTimeout(idleTimer);
       return env?.result as UpstreamCallResult;
     }
 
     if (!resp.body) {
+      clearTimeout(idleTimer);
       throw new UpstreamError(this.cfg.name, 502, "empty stream body");
     }
 
@@ -177,8 +234,24 @@ export class UpstreamAgent {
     let buf = "";
     try {
       while (true) {
-        const { done, value } = await reader.read();
-        if (value) buf += decoder.decode(value, { stream: true });
+        let done: boolean;
+        let value: Uint8Array | undefined;
+        try {
+          ({ done, value } = await reader.read());
+        } catch (err) {
+          if (timedOut) {
+            throw new UpstreamError(
+              this.cfg.name,
+              504,
+              `stream idle timeout after ${idleTimeoutMs}ms`,
+            );
+          }
+          throw err;
+        }
+        if (value) {
+          resetIdleTimer();
+          buf += decoder.decode(value, { stream: true });
+        }
         let sep = buf.indexOf("\n\n");
         while (sep !== -1) {
           const frame = buf.slice(0, sep);
@@ -189,6 +262,10 @@ export class UpstreamAgent {
           }
           if (dataLines.length > 0) {
             const raw = dataLines.join("\n");
+            if (raw === "[DONE]") {
+              sep = buf.indexOf("\n\n");
+              continue;
+            }
             let msg: any;
             try { msg = JSON.parse(raw); } catch { msg = null; }
             if (msg && msg.id === id) {
@@ -221,6 +298,7 @@ export class UpstreamAgent {
         if (done) break;
       }
     } finally {
+      clearTimeout(idleTimer);
       try { reader.releaseLock(); } catch { /* ignore */ }
     }
     throw new UpstreamError(this.cfg.name, 502, "stream ended without result");