fix(mcp-client): await failed generation shutdown
The MCP SDK starts a fire-and-forget close when initialization fails. Its stdio transport clears its process field before that close finishes, so our second Client.close() could return immediately and the reconnect timer could launch a replacement while the original child was still alive. Track the transport onclose signal for every client generation and gate failed-attempt backoff on both Client.close() settlement and that signal. Use the same barrier during plugin disposal. If the SDK's bounded stdio termination window expires without onclose, fail closed and report incomplete shutdown instead of risking overlapping server processes. Regression coverage models the SDK's early-returning second close, delayed and missing close signals, pending-connect disposal, close rejection, and the terminal timeout path. The reconnect Agent Note and Chinese counterpart now record the quiescence contract.
This commit is contained in:
@@ -44,6 +44,11 @@ export const RECONNECT_DEFAULTS: Required<ReconnectConfig> = Object.freeze({
|
||||
maxAttempts: 10,
|
||||
})
|
||||
|
||||
// The SDK's stdio transport owns two two-second termination grace periods.
|
||||
// Keep one additional second for the process-close event that proves the old
|
||||
// generation is gone; timing out fails closed instead of overlapping children.
|
||||
const GENERATION_CLOSE_TIMEOUT_MS = 5_000
|
||||
|
||||
/** Fully resolved reconnect policy captured at plugin load. */
|
||||
export type ResolvedReconnectPolicy = Readonly<Required<ReconnectConfig>>
|
||||
|
||||
@@ -133,6 +138,8 @@ export function startConnection(ctx: Context, config: Config, policy: ResolvedRe
|
||||
let disposed = false
|
||||
/** Current generation: the connecting or connected client; undefined during backoff waits and after final failure. */
|
||||
let client: Client | undefined
|
||||
/** Close signal paired with {@link client}; captured by dispose before current ownership is cleared. */
|
||||
let clientClosed: Promise<void> | undefined
|
||||
/** Live tool registrations owned by this server; only {@link enqueueSync} and dispose swap it. */
|
||||
let disposers: ToolDisposers = new Map()
|
||||
let reconnectTimer: NodeJS.Timeout | undefined
|
||||
@@ -169,9 +176,22 @@ export function startConnection(ctx: Context, config: Config, policy: ResolvedRe
|
||||
function generationDown(generation: Client): void {
|
||||
if (!isCurrent(generation)) return
|
||||
client = undefined
|
||||
clientClosed = undefined
|
||||
scheduleReconnect()
|
||||
}
|
||||
|
||||
/** Wait for the transport-owned close signal without letting a broken transport wedge teardown forever. */
|
||||
function waitForClose(closed: Promise<void>): Promise<boolean> {
|
||||
return new Promise((resolve) => {
|
||||
const timeout = setTimeout(() => { resolve(false) }, GENERATION_CLOSE_TIMEOUT_MS)
|
||||
timeout.unref()
|
||||
void closed.then(() => {
|
||||
clearTimeout(timeout)
|
||||
resolve(true)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
function scheduleReconnect(): void {
|
||||
if (!policy.enabled) {
|
||||
const detail = connectedAt !== undefined
|
||||
@@ -216,8 +236,19 @@ export function startConnection(ctx: Context, config: Config, policy: ResolvedRe
|
||||
{ name: 'dsh-mcp-client', version: '0.0.1' },
|
||||
{ capabilities: {} },
|
||||
)
|
||||
const closed: PromiseWithResolvers<void> = Promise.withResolvers()
|
||||
let attemptSettled = false
|
||||
let closeObserved = false
|
||||
const hasClosed = (): boolean => closeObserved
|
||||
client = generation
|
||||
generation.onclose = () => { generationDown(generation) }
|
||||
clientClosed = closed.promise
|
||||
generation.onclose = () => {
|
||||
closeObserved = true
|
||||
closed.resolve()
|
||||
// A failed connect owns its close barrier in the catch path below. An
|
||||
// established generation can transition down directly from this signal.
|
||||
if (attemptSettled) generationDown(generation)
|
||||
}
|
||||
// Registered before connect so a list change during the initial sync is
|
||||
// queued behind it rather than dropped.
|
||||
generation.setNotificationHandler(
|
||||
@@ -236,12 +267,32 @@ export function startConnection(ctx: Context, config: Config, policy: ResolvedRe
|
||||
)
|
||||
try {
|
||||
await generation.connect(createTransport(config))
|
||||
if (hasClosed()) {
|
||||
attemptSettled = true
|
||||
generationDown(generation)
|
||||
return
|
||||
}
|
||||
await enqueueSync(generation)
|
||||
} catch (error) {
|
||||
if (firstAttemptError === undefined) firstAttemptError = error
|
||||
// When the transport closed first, onclose already logged and scheduled.
|
||||
// Disposal clears current ownership before it closes the generation, so
|
||||
// only a live supervisor reports an attempt failure.
|
||||
if (isCurrent(generation)) ctx.logger.warn(`${label}: connection attempt failed: ${String(error)}`)
|
||||
try { await generation.close() } catch { /* transport already gone */ }
|
||||
const quiesced = hasClosed() || await waitForClose(closed.promise)
|
||||
attemptSettled = true
|
||||
if (!isCurrent(generation)) return
|
||||
if (!quiesced) {
|
||||
client = undefined
|
||||
clientClosed = undefined
|
||||
ctx.logger.error(`${label}: failed generation did not close within ${GENERATION_CLOSE_TIMEOUT_MS}ms — reconnect stopped to avoid overlapping server processes; reload the plugin or restart the Host to retry`)
|
||||
return
|
||||
}
|
||||
generationDown(generation)
|
||||
return
|
||||
}
|
||||
attemptSettled = true
|
||||
if (hasClosed()) {
|
||||
generationDown(generation)
|
||||
return
|
||||
}
|
||||
@@ -277,9 +328,14 @@ export function startConnection(ctx: Context, config: Config, policy: ResolvedRe
|
||||
reconnectTimer = undefined
|
||||
}
|
||||
const current = client
|
||||
const currentClosed = clientClosed
|
||||
client = undefined
|
||||
clientClosed = undefined
|
||||
if (current !== undefined) {
|
||||
try { await current.close() } catch { /* transport already gone */ }
|
||||
if (currentClosed !== undefined && !await waitForClose(currentClosed)) {
|
||||
ctx.logger.error(`${label}: generation did not close within ${GENERATION_CLOSE_TIMEOUT_MS}ms during disposal — server shutdown may be incomplete`)
|
||||
}
|
||||
}
|
||||
// Quiesce, don't just request it: the in-flight attempt enqueues its
|
||||
// sync before settling, so awaiting both leaves `disposers` final.
|
||||
|
||||
@@ -159,7 +159,10 @@ describe('apply (plugin lifecycle)', () => {
|
||||
beforeEach(async () => {
|
||||
vi.clearAllMocks()
|
||||
mockConnect.mockResolvedValue(undefined)
|
||||
mockClose.mockResolvedValue(undefined)
|
||||
mockClose.mockImplementation(function (this: { onclose?: () => void }) {
|
||||
this.onclose?.()
|
||||
return Promise.resolve()
|
||||
})
|
||||
mockListTools.mockResolvedValue({
|
||||
tools: [{ name: 'remote', description: 'A remote tool', inputSchema: { type: 'object' } }],
|
||||
nextCursor: undefined,
|
||||
@@ -338,7 +341,10 @@ describe('apply (plugin lifecycle)', () => {
|
||||
})
|
||||
|
||||
it('effect disposer handles client.close failure gracefully', async () => {
|
||||
mockClose.mockRejectedValue(new Error('already closed'))
|
||||
mockClose.mockImplementation(function (this: { onclose?: () => void }) {
|
||||
this.onclose?.()
|
||||
return Promise.reject(new Error('already closed'))
|
||||
})
|
||||
|
||||
await apply(ctx, stdioConfig)
|
||||
|
||||
|
||||
@@ -59,7 +59,7 @@ vi.mock('@modelcontextprotocol/sdk/client/streamableHttp.js', () => ({
|
||||
// vi.mock is hoisted above static imports, so the modules under test see the
|
||||
// mocked SDK even through a static import.
|
||||
import { apply } from '@deepseek-ai/dsh-mcp-client/src/index.ts'
|
||||
import { RECONNECT_DEFAULTS, resolveReconnectPolicy } from '@deepseek-ai/dsh-mcp-client/src/connection.ts'
|
||||
import { RECONNECT_DEFAULTS, resolveReconnectPolicy, startConnection } from '@deepseek-ai/dsh-mcp-client/src/connection.ts'
|
||||
|
||||
// ---- Helpers ----
|
||||
|
||||
@@ -128,7 +128,10 @@ describe('reconnect supervisor', () => {
|
||||
vi.clearAllMocks()
|
||||
instances.length = 0
|
||||
mockConnect.mockResolvedValue(undefined)
|
||||
mockClose.mockResolvedValue(undefined)
|
||||
mockClose.mockImplementation(function (this: { onclose?: () => void }) {
|
||||
this.onclose?.()
|
||||
return Promise.resolve()
|
||||
})
|
||||
mockListTools.mockResolvedValue(listing('remote'))
|
||||
mockCallTool.mockResolvedValue({ content: [{ type: 'text', text: 'ok' }] })
|
||||
ctx = await mountRegistry()
|
||||
@@ -174,7 +177,10 @@ describe('reconnect supervisor', () => {
|
||||
|
||||
mockConnect.mockRejectedValue(new Error('server gone'))
|
||||
// A failing close on the failed attempt's cleanup must not break the loop.
|
||||
mockClose.mockRejectedValue(new Error('already closed'))
|
||||
mockClose.mockImplementation(function (this: { onclose?: () => void }) {
|
||||
this.onclose?.()
|
||||
return Promise.reject(new Error('already closed'))
|
||||
})
|
||||
instances[0]!.onclose?.()
|
||||
|
||||
await vi.waitFor(() => {
|
||||
@@ -189,6 +195,78 @@ describe('reconnect supervisor', () => {
|
||||
expect(mockConnect).toHaveBeenCalledTimes(3)
|
||||
})
|
||||
|
||||
it('does not start a replacement until a failed generation reports that it closed', async () => {
|
||||
mockConnect.mockRejectedValueOnce(new Error('initialize failed'))
|
||||
// Model the SDK's fire-and-forget close after initialize fails: the
|
||||
// harness's second close call returns, but the child has not exited yet.
|
||||
mockClose.mockResolvedValue(undefined)
|
||||
|
||||
const applying = apply(ctx, stdioConfig({ initialDelayMs: 2, maxDelayMs: 8, maxAttempts: 2 }))
|
||||
await vi.waitFor(() => { expect(mockClose).toHaveBeenCalled() })
|
||||
await sleep(30)
|
||||
expect(instances).toHaveLength(1)
|
||||
|
||||
instances[0]!.onclose?.()
|
||||
await applying
|
||||
await vi.waitFor(() => { expect(instances).toHaveLength(2) })
|
||||
})
|
||||
|
||||
it('stops reconnecting when a failed generation never reports that it closed', async () => {
|
||||
vi.useFakeTimers()
|
||||
try {
|
||||
const { errors } = captureLogs(ctx)
|
||||
mockConnect.mockRejectedValue(new Error('initialize failed'))
|
||||
mockClose.mockResolvedValue(undefined)
|
||||
|
||||
const applying = apply(ctx, stdioConfig({ initialDelayMs: 2, maxDelayMs: 8, maxAttempts: 2 }))
|
||||
await vi.advanceTimersByTimeAsync(5_000)
|
||||
await applying
|
||||
|
||||
expect(instances).toHaveLength(1)
|
||||
expect(errors.some(line => line.includes('reconnect stopped to avoid overlapping server processes'))).toBe(true)
|
||||
} finally {
|
||||
vi.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
it('suppresses retry reporting when disposal owns a pending connect rejection', async () => {
|
||||
const { warns } = captureLogs(ctx)
|
||||
const gate: PromiseWithResolvers<void> = Promise.withResolvers()
|
||||
mockConnect.mockImplementation(() => gate.promise)
|
||||
const handle = startConnection(ctx, stdioConfig(), resolveReconnectPolicy(undefined, 'reconnect'))
|
||||
await vi.waitFor(() => { expect(instances).toHaveLength(1) })
|
||||
|
||||
const disposing = handle.dispose()
|
||||
gate.reject(new Error('disposed connect'))
|
||||
await disposing
|
||||
await handle.ready
|
||||
|
||||
expect(warns.some(line => line.includes('connection attempt failed'))).toBe(false)
|
||||
expect(instances).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('bounds disposal while a resolving generation never reports that it closed', async () => {
|
||||
vi.useFakeTimers()
|
||||
try {
|
||||
const { errors } = captureLogs(ctx)
|
||||
const gate: PromiseWithResolvers<void> = Promise.withResolvers()
|
||||
mockConnect.mockImplementation(() => gate.promise)
|
||||
mockClose.mockResolvedValue(undefined)
|
||||
const handle = startConnection(ctx, stdioConfig(), resolveReconnectPolicy(undefined, 'reconnect'))
|
||||
await vi.advanceTimersByTimeAsync(0)
|
||||
|
||||
const disposing = handle.dispose()
|
||||
await vi.advanceTimersByTimeAsync(5_000)
|
||||
gate.resolve()
|
||||
await disposing
|
||||
|
||||
expect(mockListTools).not.toHaveBeenCalled()
|
||||
expect(errors.some(line => line.includes('server shutdown may be incomplete'))).toBe(true)
|
||||
} finally {
|
||||
vi.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
it('dispose during the backoff wait cancels the pending reconnect', async () => {
|
||||
await apply(ctx, stdioConfig({ initialDelayMs: 60_000, maxDelayMs: 60_000, maxAttempts: 5 }))
|
||||
await vi.waitFor(() => { expect(ctx.tools.get('mcp__srv__remote')).toBeDefined() })
|
||||
|
||||
Reference in New Issue
Block a user