diff --git a/src/main.ts b/src/main.ts index 40c241b47..3d32798b7 100644 --- a/src/main.ts +++ b/src/main.ts @@ -82,6 +82,7 @@ export const run = async ( createShutdown({ handle: transport.handle, closeAndroid: transport.closeAndroid, + closeOptionalProviders: transport.closeOptionalProviders, session, dependencies, serverLogger, diff --git a/src/main/shutdown.ts b/src/main/shutdown.ts index f6722eb97..90d30c575 100644 --- a/src/main/shutdown.ts +++ b/src/main/shutdown.ts @@ -16,6 +16,7 @@ interface ShutdownCleanupOwner { export const createShutdown = (input: { readonly handle: StdioServerHandle; readonly closeAndroid?: () => Promise; + readonly closeOptionalProviders?: () => Promise; readonly session: BinarySession; readonly dependencies: RuntimeDependencies; readonly serverLogger: Logger; @@ -30,6 +31,9 @@ export const createShutdown = (input: { ...(input.closeAndroid === undefined ? [] : [{ close: input.closeAndroid, closed: false }]), + ...(input.closeOptionalProviders === undefined + ? [] + : [{ close: input.closeOptionalProviders, closed: false }]), { close: async () => { const closed = await session.close({ retainProviderDocuments: true }); diff --git a/src/main/transport.ts b/src/main/transport.ts index f8d6014f3..f9a5cd0ea 100644 --- a/src/main/transport.ts +++ b/src/main/transport.ts @@ -31,6 +31,7 @@ export const startMcpTransport = async ( readonly ok: true; readonly handle: StdioServerHandle; readonly closeAndroid: () => Promise; + readonly closeOptionalProviders: () => Promise; } | { readonly ok: false } > => { @@ -62,6 +63,11 @@ export const startMcpTransport = async ( for (const result of results) if (result.status === "rejected") throw result.reason; }; + // Optional providers are shared by every server instance the SDK creates, + // including a discarded discovery probe, so only process shutdown closes them. + const closeOptionalProviders = async (): Promise => { + await optionalProviders.browserScenarioCapture?.close?.(); + }; let handle: StdioServerHandle; try { handle = dependencies.serve( @@ -124,5 +130,5 @@ export const startMcpTransport = async ( } return { ok: false }; } - return { ok: true, handle, closeAndroid }; + return { ok: true, handle, closeAndroid, closeOptionalProviders }; }; diff --git a/src/server/createServer.ts b/src/server/createServer.ts index 0e5c69cba..3220fe8e7 100644 --- a/src/server/createServer.ts +++ b/src/server/createServer.ts @@ -270,7 +270,6 @@ export const createServer = ( javascriptRecovery.close?.() ?? Promise.resolve(), artifactResources.close(), firmwareAnalysis.close(), - options.browserScenarioCapture?.close?.() ?? Promise.resolve(), ]) .then((results) => { const failures = results.flatMap((result) => diff --git a/tests/boundary/process/browserScenarioProbeLifecycle.test.ts b/tests/boundary/process/browserScenarioProbeLifecycle.test.ts new file mode 100644 index 000000000..42871d13a --- /dev/null +++ b/tests/boundary/process/browserScenarioProbeLifecycle.test.ts @@ -0,0 +1,61 @@ +import { STDIO_DEFAULT_MAX_BUFFER_SIZE } from "@modelcontextprotocol/server"; +import { expect, it, vi } from "vitest"; + +import type { BrowserScenarioCapturePort } from "../../../src/application/BrowserScenarioCapturePort.js"; +import { silentLogger } from "../../../src/logger.js"; +import { startMcpTransport } from "../../../src/main/transport.js"; +import type { RuntimeDependencies } from "../../../src/main/types.js"; +import { ToolResultDelivery } from "../../../src/server/toolResult.js"; +import { createTestBinarySession } from "../../fixtures/binarySession.js"; + +it("keeps the shared browser scenario provider open when the SDK discards its discovery probe", async () => { + const session = createTestBinarySession(() => { + throw new Error("Unexpected deep provider"); + }); + const close = vi.fn(async () => undefined); + const browserScenarioCapture: BrowserScenarioCapturePort = { + identity: () => ({ + id: "fake-browser", + name: "Fake browser", + version: "1", + }), + captureScenario: async () => { + throw new Error("Unexpected scenario capture"); + }, + close, + }; + const factories: Parameters[0][] = []; + const runtime: RuntimeDependencies = { + env: {}, + serve: (factory) => { + factories.push(factory); + return { close: async () => undefined }; + }, + writeStderr: () => undefined, + setExitCode: () => undefined, + registerShutdown: () => () => undefined, + }; + const transport = await startMcpTransport(runtime, session, { + logger: silentLogger, + serverLogger: silentLogger, + environment: runtime.env, + delivery: new ToolResultDelivery(STDIO_DEFAULT_MAX_BUFFER_SIZE), + loadOptionalProviders: async () => ({ browserScenarioCapture }), + }); + if (!transport.ok) throw new Error("Transport failed"); + const factory = factories[0]; + if (factory === undefined) throw new Error("Missing SDK server factory"); + try { + const probe = await factory({ era: "modern" }); + await probe.close(); + const replacement = await factory({ era: "legacy" }); + await replacement.close(); + expect(close).not.toHaveBeenCalled(); + + await transport.closeOptionalProviders(); + expect(close).toHaveBeenCalledTimes(1); + } finally { + await transport.closeAndroid(); + await session.close(); + } +});