All files / src/mcp/transport reconnect.ts

100% Statements 39/39
85.71% Branches 12/14
100% Functions 6/6
100% Lines 34/34

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162                                                                                                                            7x 7x 1x   6x 6x 6x 6x 6x 6x   6x                     6x 6x 6x       6x 6x 6x 6x   1x                                             9x 7x 7x       2x   9x 1x   9x                                     186x 186x 195x 195x   86x 86x 11x 9x     2x    
// SPDX-FileCopyrightText: 2024-2026 Hack23 AB
// SPDX-License-Identifier: Apache-2.0
 
/**
 * @module MCP/transport/reconnect
 * @description Exponential back-off reconnect loop and tool-call retry logic
 * for MCPConnection.
 *
 * Extracted from `connection.ts` to keep individual file sizes under 400 LOC.
 * Operates on an explicit {@link ReconnectOps} adapter rather than `this`.
 */
 
import type { MCPToolResult } from '../../types/index.js';
import { isRetriableError, RECONNECT_MAX_DELAY_MS } from './retry-policy.js';
 
// ─── Context interface ────────────────────────────────────────────────────────
 
/**
 * Adapter passed by {@link MCPConnection} to reconnect/retry helpers.
 * Exposes the handful of connection fields these helpers need through
 * getter/setter callbacks so the fields remain on the connection class
 * (preserving external observability via `getConnectionHealth()`).
 */
export interface ReconnectOps {
  /** Maximum consecutive connection attempts before giving up */
  readonly maxConnectionAttempts: number;
  /** Base delay (ms) between retry attempts */
  readonly connectionRetryDelay: number;
  /** Human-readable label for log messages */
  readonly serverLabel: string;
  /** Call the underlying MCP tool (used by the retry loop) */
  readonly callTool: (name: string, args: object) => Promise<MCPToolResult>;
  /** Whether the transport is currently connected */
  readonly isConnected: () => boolean;
  /** Trigger a full reconnect cycle */
  readonly connect: () => Promise<void>;
  /** Mark the transport as (dis)connected */
  readonly setConnected: (v: boolean) => void;
  /** Current reconnect-attempt counter */
  readonly getReconnectCount: () => number;
  /** Update the reconnect-attempt counter */
  readonly setReconnectCount: (n: number) => void;
  /** Current in-flight reconnect promise (or null) */
  readonly getReconnectingPromise: () => Promise<void> | null;
  /** Persist or clear the in-flight reconnect promise */
  readonly setReconnectingPromise: (p: Promise<void> | null) => void;
  /** Current cumulative timeout counter */
  readonly getTimeoutCount: () => number;
  /** Update the cumulative timeout counter */
  readonly setTimeoutCount: (n: number) => void;
}
 
// ─── Exported helpers ─────────────────────────────────────────────────────────
 
/**
 * Reconnect with exponential back-off. Concurrent callers await the same
 * in-flight reconnect promise instead of spawning parallel attempts.
 *
 * @param ops - Reconnect operations adapter from MCPConnection
 * @returns Promise that resolves when reconnection succeeds or all attempts are exhausted
 */
export async function performReconnect(ops: ReconnectOps): Promise<void> {
  const inflight = ops.getReconnectingPromise();
  if (inflight !== null) {
    return inflight;
  }
  ops.setReconnectCount(ops.getReconnectCount() + 1);
  console.log(`🔄 Reconnecting to ${ops.serverLabel} (attempt ${ops.getReconnectCount()})...`);
  const p = doReconnect(ops);
  ops.setReconnectingPromise(p);
  try {
    await p;
  } finally {
    ops.setReconnectingPromise(null);
  }
}
 
/**
 * Internal reconnect helper. Waits for an exponential back-off delay then
 * delegates to `connect()`, which handles its own retry loop.
 *
 * @param ops - Reconnect operations adapter
 */
async function doReconnect(ops: ReconnectOps): Promise<void> {
  const normalizedMax = Math.max(1, ops.maxConnectionAttempts);
  const attemptIndex = Math.min(Math.max(0, ops.getReconnectCount() - 1), normalizedMax - 1);
  const delay = Math.min(
    ops.connectionRetryDelay * Math.pow(2, attemptIndex),
    RECONNECT_MAX_DELAY_MS
  );
  await new Promise((r) => setTimeout(r, delay));
  try {
    ops.setConnected(false);
    await ops.connect();
  } catch (error) {
    console.error(
      `❌ Reconnection to ${ops.serverLabel} failed: ${
        error instanceof Error ? error.message : String(error)
      }`
    );
  }
}
 
/**
 * Log a retry warning and, if disconnected, attempt to reconnect before waiting.
 *
 * @param lastError - The error from the failed attempt
 * @param attempt - Zero-based current attempt index
 * @param retries - Total retry count
 * @param ops - Reconnect operations adapter
 * @returns Promise that resolves after logging, optional reconnect, and inter-retry delay
 */
export async function handleRetryAttempt(
  lastError: Error,
  attempt: number,
  retries: number,
  ops: ReconnectOps
): Promise<void> {
  if (lastError.message.toLowerCase().includes('timeout')) {
    ops.setTimeoutCount(ops.getTimeoutCount() + 1);
    console.warn(
      `⏱️ Request timeout (total: ${ops.getTimeoutCount()}), retrying ${attempt + 1}/${retries}...`
    );
  } else {
    console.warn(`⚠️ Request failed, retrying ${attempt + 1}/${retries}: ${lastError.message}`);
  }
  if (!ops.isConnected()) {
    await performReconnect(ops);
  }
  await new Promise((r) => setTimeout(r, ops.connectionRetryDelay * (attempt + 1)));
}
 
/**
 * Call an MCP tool with automatic retry on timeout or connection loss.
 * Non-retriable errors are re-thrown immediately without consuming retry budget.
 *
 * @param name - Tool name
 * @param args - Tool arguments (plain object, non-null, not an array)
 * @param retries - Maximum number of retries (validated ≥ 0 by caller)
 * @param ops - Reconnect operations adapter
 * @returns Tool execution result
 */
export async function runWithRetry(
  name: string,
  args: object,
  retries: number,
  ops: ReconnectOps
): Promise<MCPToolResult> {
  let lastError: Error = new Error(`Failed to call tool '${name}' after ${retries} retries`);
  for (let attempt = 0; attempt <= retries; attempt++) {
    try {
      return await ops.callTool(name, args);
    } catch (error) {
      lastError = error instanceof Error ? error : new Error(String(error));
      if (!isRetriableError(lastError)) throw lastError;
      if (attempt === retries) break;
      await handleRetryAttempt(lastError, attempt, retries, ops);
    }
  }
  throw lastError;
}