600 lines
17 KiB
TypeScript
600 lines
17 KiB
TypeScript
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
|
import { ResilientTransportWrapper } from '../src/resilient.js';
|
|
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
|
|
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
|
|
|
|
class MockTransport implements Transport {
|
|
onclose?: () => void;
|
|
onerror?: (error: Error) => void;
|
|
onmessage?: (message: JSONRPCMessage) => void;
|
|
|
|
startCalls = 0;
|
|
closeCalls = 0;
|
|
sentMessages: JSONRPCMessage[] = [];
|
|
shouldFailStart = false;
|
|
shouldFailSend = false;
|
|
|
|
async start(): Promise<void> {
|
|
this.startCalls++;
|
|
if (this.shouldFailStart) {
|
|
throw new Error('Mock start failure');
|
|
}
|
|
}
|
|
|
|
async close(): Promise<void> {
|
|
this.closeCalls++;
|
|
if (this.onclose) {
|
|
this.onclose();
|
|
}
|
|
}
|
|
|
|
async send(message: JSONRPCMessage): Promise<void> {
|
|
if (this.shouldFailSend) {
|
|
throw new Error('Mock send failure');
|
|
}
|
|
this.sentMessages.push(message);
|
|
}
|
|
|
|
// Helper to simulate incoming messages
|
|
simulateMessage(msg: JSONRPCMessage) {
|
|
if (this.onmessage) {
|
|
this.onmessage(msg);
|
|
}
|
|
}
|
|
|
|
// Helper to simulate transport error
|
|
simulateError(err: Error) {
|
|
if (this.onerror) {
|
|
this.onerror(err);
|
|
}
|
|
}
|
|
|
|
// Helper to simulate transport close
|
|
simulateClose() {
|
|
if (this.onclose) {
|
|
this.onclose();
|
|
}
|
|
}
|
|
}
|
|
|
|
describe('ResilientTransportWrapper', () => {
|
|
let mockTransports: MockTransport[] = [];
|
|
|
|
beforeEach(() => {
|
|
mockTransports = [];
|
|
vi.useFakeTimers();
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.useRealTimers();
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
const createWrapper = (options = {}) => {
|
|
return new ResilientTransportWrapper({
|
|
name: 'test-wrapper',
|
|
connectParams: async () => {
|
|
const t = new MockTransport();
|
|
mockTransports.push(t);
|
|
return t;
|
|
},
|
|
...options
|
|
});
|
|
};
|
|
|
|
it('should connect successfully', async () => {
|
|
const wrapper = createWrapper();
|
|
await wrapper.start();
|
|
|
|
expect(mockTransports.length).toBe(1);
|
|
expect(mockTransports[0].startCalls).toBe(1);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should queue messages while reconnecting, then flush them', async () => {
|
|
const wrapper = createWrapper({ reconnectDelayMs: 100 });
|
|
await wrapper.start();
|
|
|
|
const firstTransport = mockTransports[0];
|
|
|
|
// Simulate drop
|
|
firstTransport.simulateClose();
|
|
|
|
// Wrapper state should now be reconnecting
|
|
|
|
const testMsg: JSONRPCMessage = { jsonrpc: '2.0', method: 'test', params: { a: 1 } };
|
|
|
|
// Send while reconnecting
|
|
await wrapper.send(testMsg);
|
|
|
|
// Ensure it wasn't sent to the closed transport
|
|
expect(firstTransport.sentMessages.length).toBe(0);
|
|
|
|
// Fast forward to trigger reconnect
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
|
|
// A new transport should be created
|
|
expect(mockTransports.length).toBe(2);
|
|
const secondTransport = mockTransports[1];
|
|
|
|
// The queued message should have been flushed
|
|
expect(secondTransport.sentMessages.length).toBe(1);
|
|
expect(secondTransport.sentMessages[0]).toEqual(testMsg);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should reconnect if heartbeat fails repeatedly', async () => {
|
|
let healthCheckCalls = 0;
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
pingTimeoutMs: 500,
|
|
maxConsecutiveFailures: 2,
|
|
reconnectDelayMs: 100,
|
|
healthCheckFn: async () => {
|
|
healthCheckCalls++;
|
|
return false; // Always fail health check
|
|
}
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
|
|
// Trigger first heartbeat failure
|
|
await vi.advanceTimersByTimeAsync(1001);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
expect(mockTransports.length).toBe(1); // Still same transport
|
|
|
|
// Trigger second heartbeat failure (maxConsecutiveFailures = 2)
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
|
|
// Advance for reconnect delay
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
|
|
// Should have reconnected (new transport created)
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should intercept inner transport errors and trigger reconnect', async () => {
|
|
const wrapper = createWrapper({ reconnectDelayMs: 100 });
|
|
await wrapper.start();
|
|
|
|
const firstTransport = mockTransports[0];
|
|
firstTransport.simulateError(new Error('Test error'));
|
|
|
|
// Advance time for reconnect delay
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
|
|
// Should have reconnected
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should pass transparent messages down to onmessage handler', async () => {
|
|
const wrapper = createWrapper();
|
|
await wrapper.start();
|
|
|
|
let receivedMsg: JSONRPCMessage | null = null;
|
|
wrapper.onmessage = (msg) => {
|
|
receivedMsg = msg;
|
|
};
|
|
|
|
const testMsg: JSONRPCMessage = { jsonrpc: '2.0', method: 'notify', params: { x: 1 } };
|
|
mockTransports[0].simulateMessage(testMsg);
|
|
|
|
expect(receivedMsg).toEqual(testMsg);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should handle flushQueue errors by putting messages back', async () => {
|
|
const wrapper = createWrapper({ maxQueueSize: 5 });
|
|
await wrapper.start();
|
|
|
|
const firstTransport = mockTransports[0];
|
|
firstTransport.simulateClose(); // puts it in reconnecting state
|
|
|
|
// Add messages to queue
|
|
await wrapper.send({ jsonrpc: '2.0', method: 'test1' });
|
|
await wrapper.send({ jsonrpc: '2.0', method: 'test2' });
|
|
|
|
// Advance time to reconnect
|
|
let newTransportReady = false;
|
|
const originalSetTimeout = global.setTimeout;
|
|
const spy = vi.spyOn(global, 'setTimeout').mockImplementationOnce((cb: any, ms?: number) => {
|
|
// simulate the factory creating the second one and failing its send
|
|
const id = originalSetTimeout(() => {
|
|
cb();
|
|
if (mockTransports.length > 1) {
|
|
mockTransports[1].shouldFailSend = true;
|
|
newTransportReady = true;
|
|
}
|
|
}, ms);
|
|
return id as any;
|
|
});
|
|
|
|
await vi.advanceTimersByTimeAsync(3000);
|
|
|
|
// Test the queue size management limit
|
|
firstTransport.simulateClose();
|
|
for (let i = 0; i < 10; i++) {
|
|
await wrapper.send({ jsonrpc: '2.0', method: `spam${i}` });
|
|
}
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should throw when sending on a closed transport', async () => {
|
|
const wrapper = createWrapper();
|
|
await wrapper.start();
|
|
await wrapper.close();
|
|
|
|
await expect(wrapper.send({ jsonrpc: '2.0', method: 'test' })).rejects.toThrow('Transport is closed');
|
|
});
|
|
|
|
it('should throw when initial connect fails', async () => {
|
|
const wrapper = new ResilientTransportWrapper({
|
|
name: 'fail-connect',
|
|
reconnectDelayMs: 100,
|
|
connectParams: async () => {
|
|
throw new Error('Initial network error');
|
|
}
|
|
});
|
|
|
|
// start() should throw on initial connection failure
|
|
// Caller is responsible for retry logic
|
|
await expect(wrapper.start()).rejects.toThrow('Initial network error');
|
|
|
|
// Clean up
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should call healthCheckFn for heartbeat', async () => {
|
|
let healthCheckCalls = 0;
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
pingTimeoutMs: 500,
|
|
maxConsecutiveFailures: 3,
|
|
reconnectDelayMs: 100,
|
|
healthCheckFn: async () => {
|
|
healthCheckCalls++;
|
|
return true;
|
|
}
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
|
|
// Advance timer to trigger first heartbeat
|
|
await vi.advanceTimersByTimeAsync(1001);
|
|
// Allow async healthCheckFn to complete
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
|
|
// healthCheckFn should have been called once
|
|
expect(healthCheckCalls).toBe(1);
|
|
|
|
// Advance for another heartbeat
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
expect(healthCheckCalls).toBe(2);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should trigger reconnect when healthCheckFn fails repeatedly', async () => {
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
pingTimeoutMs: 500,
|
|
maxConsecutiveFailures: 2,
|
|
reconnectDelayMs: 100,
|
|
healthCheckFn: async () => {
|
|
return false; // Always fail
|
|
}
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
|
|
// Trigger first heartbeat failure
|
|
await vi.advanceTimersByTimeAsync(1001);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
expect(mockTransports.length).toBe(1); // Still same transport
|
|
|
|
// Trigger second heartbeat failure (maxConsecutiveFailures = 2)
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
|
|
// Advance for reconnect delay
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
|
|
// Should have reconnected (new transport created)
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should trigger onreconnect event on reconnect', async () => {
|
|
let onreconnectCalls = 0;
|
|
const wrapper = createWrapper({ reconnectDelayMs: 100 });
|
|
|
|
wrapper.onreconnect = async () => {
|
|
onreconnectCalls++;
|
|
};
|
|
|
|
await wrapper.start();
|
|
|
|
// Simulate disconnect
|
|
mockTransports[0].simulateClose();
|
|
|
|
// Advance time for reconnect
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
|
|
// onreconnect should have been called
|
|
expect(onreconnectCalls).toBe(1);
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should handle onreconnect failure with backoff', async () => {
|
|
let reconnectAttempts = 0;
|
|
const wrapper = createWrapper({ reconnectDelayMs: 100 });
|
|
|
|
wrapper.onreconnect = async () => {
|
|
reconnectAttempts++;
|
|
throw new Error('Reconnect handler failed');
|
|
};
|
|
|
|
await wrapper.start();
|
|
|
|
// Simulate disconnect
|
|
mockTransports[0].simulateClose();
|
|
|
|
// Advance time for first reconnect attempt (100ms delay)
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
expect(reconnectAttempts).toBe(1);
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
// onreconnect failed, should retry with backoff (400ms = 100 * 2^2)
|
|
await vi.advanceTimersByTimeAsync(450);
|
|
expect(reconnectAttempts).toBe(2);
|
|
expect(mockTransports.length).toBe(3);
|
|
|
|
// Next retry with backoff (800ms = 100 * 2^3)
|
|
await vi.advanceTimersByTimeAsync(850);
|
|
expect(reconnectAttempts).toBe(3);
|
|
expect(mockTransports.length).toBe(4);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should succeed when onreconnect handler succeeds', async () => {
|
|
let reconnectAttempts = 0;
|
|
const wrapper = createWrapper({ reconnectDelayMs: 100 });
|
|
|
|
wrapper.onreconnect = async () => {
|
|
reconnectAttempts++;
|
|
// Success - no throw
|
|
};
|
|
|
|
await wrapper.start();
|
|
|
|
// Simulate disconnect
|
|
mockTransports[0].simulateClose();
|
|
|
|
// Advance time for reconnect
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
|
|
// onreconnect should have succeeded
|
|
expect(reconnectAttempts).toBe(1);
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
// No more reconnects should happen
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
expect(reconnectAttempts).toBe(1);
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should use transport liveness check when no healthCheckFn provided', async () => {
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
pingTimeoutMs: 500,
|
|
maxConsecutiveFailures: 2,
|
|
reconnectDelayMs: 100,
|
|
// No healthCheckFn provided
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
|
|
// Trigger heartbeat - should pass because transport is alive
|
|
await vi.advanceTimersByTimeAsync(1001);
|
|
|
|
// Should still be on same transport (health check passed via liveness)
|
|
expect(mockTransports.length).toBe(1);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should prevent concurrent health check executions', async () => {
|
|
let healthCheckCalls = 0;
|
|
let healthCheckResolve: () => void;
|
|
let healthCheckPromise: Promise<boolean>;
|
|
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
pingTimeoutMs: 5000, // Long timeout to test concurrency
|
|
maxConsecutiveFailures: 3,
|
|
healthCheckFn: async () => {
|
|
healthCheckCalls++;
|
|
// Return a promise that doesn't resolve immediately
|
|
healthCheckPromise = new Promise((resolve) => {
|
|
healthCheckResolve = () => resolve(true);
|
|
});
|
|
return healthCheckPromise;
|
|
}
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
|
|
// Trigger first heartbeat - this will start but not complete
|
|
await vi.advanceTimersByTimeAsync(1001);
|
|
|
|
// healthCheckFn should have been called once and is now pending
|
|
expect(healthCheckCalls).toBe(1);
|
|
|
|
// Trigger more timers while health check is in progress
|
|
// These should be skipped due to healthCheckInProgress guard
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
|
|
// Still only 1 call because the first one hasn't resolved yet
|
|
expect(healthCheckCalls).toBe(1);
|
|
|
|
// Now resolve the health check
|
|
healthCheckResolve!();
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
|
|
// Now the next heartbeat can run
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
expect(healthCheckCalls).toBe(2);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should use response-driven scheduling (setTimeout not setInterval)', async () => {
|
|
let healthCheckCalls = 0;
|
|
const healthCheckTimes: number[] = [];
|
|
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
healthCheckFn: async () => {
|
|
healthCheckCalls++;
|
|
healthCheckTimes.push(Date.now());
|
|
return true;
|
|
}
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
|
|
// Run for 5 heartbeat cycles
|
|
for (let i = 0; i < 5; i++) {
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10); // Allow async to complete
|
|
}
|
|
|
|
expect(healthCheckCalls).toBe(5);
|
|
|
|
// Verify timing - each check should be ~1000ms apart
|
|
// (not stacked up like setInterval would do if checks were slow)
|
|
for (let i = 1; i < healthCheckTimes.length; i++) {
|
|
const diff = healthCheckTimes[i] - healthCheckTimes[i - 1];
|
|
expect(diff).toBeGreaterThanOrEqual(1000);
|
|
}
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should not schedule next health check when state changes to reconnecting', async () => {
|
|
let healthCheckCalls = 0;
|
|
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
maxConsecutiveFailures: 2,
|
|
reconnectDelayMs: 100,
|
|
healthCheckFn: async () => {
|
|
healthCheckCalls++;
|
|
return false; // Always fail
|
|
}
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
|
|
// Trigger first failure
|
|
await vi.advanceTimersByTimeAsync(1001);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
expect(healthCheckCalls).toBe(1);
|
|
|
|
// Trigger second failure - should trigger reconnect
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
expect(healthCheckCalls).toBe(2);
|
|
|
|
// After hitting maxConsecutiveFailures, no more health checks should be scheduled
|
|
// because state is now 'reconnecting'
|
|
const callsBeforeReconnect = healthCheckCalls;
|
|
|
|
// Advance time - no new health checks should run during reconnect
|
|
await vi.advanceTimersByTimeAsync(50);
|
|
expect(healthCheckCalls).toBe(callsBeforeReconnect);
|
|
|
|
// Complete the reconnect
|
|
await vi.advanceTimersByTimeAsync(100);
|
|
|
|
// After reconnect, heartbeat should restart
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
expect(healthCheckCalls).toBeGreaterThan(callsBeforeReconnect);
|
|
|
|
await wrapper.close();
|
|
});
|
|
|
|
it('should schedule next check in finally block only when connected', async () => {
|
|
let healthCheckCalls = 0;
|
|
let shouldFailHealthCheck = true;
|
|
|
|
const wrapper = createWrapper({
|
|
pingIntervalMs: 1000,
|
|
maxConsecutiveFailures: 2,
|
|
reconnectDelayMs: 100,
|
|
healthCheckFn: async () => {
|
|
healthCheckCalls++;
|
|
if (shouldFailHealthCheck) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
});
|
|
|
|
await wrapper.start();
|
|
wrapper.enableHeartbeat();
|
|
|
|
// Trigger failures to cause reconnect
|
|
await vi.advanceTimersByTimeAsync(1001);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
|
|
// Now let health checks pass
|
|
shouldFailHealthCheck = false;
|
|
|
|
// Complete reconnect
|
|
await vi.advanceTimersByTimeAsync(150);
|
|
|
|
// Should be reconnected now
|
|
expect(mockTransports.length).toBe(2);
|
|
|
|
// Heartbeat should resume with new transport
|
|
const callsAfterReconnect = healthCheckCalls;
|
|
await vi.advanceTimersByTimeAsync(1000);
|
|
await vi.advanceTimersByTimeAsync(10);
|
|
|
|
// New health check should have been scheduled after reconnect
|
|
expect(healthCheckCalls).toBeGreaterThan(callsAfterReconnect);
|
|
|
|
await wrapper.close();
|
|
});
|
|
});
|