diff --git a/packages/server/src/server/plugins/lifecycle.e2e.test.ts b/packages/server/src/server/plugins/lifecycle.e2e.test.ts index 05157e3816..aeaa1c79bc 100644 --- a/packages/server/src/server/plugins/lifecycle.e2e.test.ts +++ b/packages/server/src/server/plugins/lifecycle.e2e.test.ts @@ -1,6 +1,7 @@ -import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { cp, mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import path from "node:path"; +import { fileURLToPath } from "node:url"; import { expect, onTestFinished, test } from "vitest"; import { DaemonClient } from "../test-utils/daemon-client.js"; import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js"; @@ -457,3 +458,80 @@ test("a provider that never acknowledges session.close does not hold the daemon // test fails by timing out. await expect(daemon.daemon.stop()).resolves.toBeUndefined(); }, 90_000); + +const repoRoot = fileURLToPath(new URL("../../../../../", import.meta.url)); + +// Hold the provider close until the subprocess has entered shutdown cleanup. +async function createSlowClosingProviderPlugin() { + const directory = await mkdtemp(path.join(tmpdir(), "paseo-slow-close-plugin-")); + const closeStarted = path.join(directory, "close-started"); + const closeRelease = path.join(directory, "close-release"); + const cleanupStarted = path.join(directory, "cleanup-started"); + await cp(path.join(repoRoot, "plugin-examples/provider-direct"), directory, { recursive: true }); + const providerPath = path.join(directory, "server", "provider.ts"); + const source = (await readFile(providerPath, "utf8")).replaceAll("\r\n", "\n"); + const patched = `import { readFile, writeFile } from "node:fs/promises";\n${source.replace( + " closed = true;\n sessions.clear();", + ` closed = true; + await writeFile(${JSON.stringify(closeStarted)}, "started"); + while (true) { + try { await readFile(${JSON.stringify(closeRelease)}); break; } + catch { await new Promise((resolve) => setTimeout(resolve, 10)); } + } + sessions.clear(); + console.log("provider close completed");`, + )}`; + expect(patched).toContain(`await writeFile(${JSON.stringify(closeStarted)}, "started")`); + await writeFile(providerPath, patched); + const entryPath = path.join(directory, "index.server.ts"); + const entry = (await readFile(entryPath, "utf8")).replaceAll("\r\n", "\n"); + const withCleanup = `import { writeFile } from "node:fs/promises";\n${entry.replace( + " return () => {};", + ` return async () => { await writeFile(${JSON.stringify(cleanupStarted)}, "started"); console.log("plugin cleanup completed"); };`, + )}`; + expect(withCleanup).toContain(`await writeFile(${JSON.stringify(cleanupStarted)}, "started")`); + await writeFile(entryPath, withCleanup); + return { directory, closeStarted, closeRelease, cleanupStarted }; +} + +test("reloading a plugin does not send on a closed IPC channel", async () => { + const { + directory: pluginDirectory, + closeStarted, + closeRelease, + cleanupStarted, + } = await createSlowClosingProviderPlugin(); + const cwd = await mkdtemp(path.join(tmpdir(), "paseo-reload-teardown-")); + const daemon = await createTestPaseoDaemon(); + const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + try { + await client.connect(); + await client.patchDaemonConfig({ pluginsEnabled: true }); + await client.installDirectoryPlugin(pluginDirectory); + const agent = await client.createAgent({ provider: "direct-example", cwd }); + const turn = client.sendMessage(agent.id, "Say hello").catch(() => undefined); + const reload = client.reloadPlugin("provider-direct-example"); + await expect + .poll(() => readFile(closeStarted, "utf8").catch(() => null), { timeout: 20_000 }) + .toBe("started"); + await expect + .poll(() => readFile(cleanupStarted, "utf8").catch(() => null), { timeout: 20_000 }) + .toBe("started"); + await writeFile(closeRelease, "release"); + const reloaded = await reload; + expect(reloaded.status).toBe("running"); + const entries = await client.getPluginLogs("provider-direct-example"); + const messages = entries.map((entry) => entry.message).join("\n"); + expect(messages).not.toContain("ERR_IPC_CHANNEL_CLOSED"); + expect(entries.filter((entry) => entry.message === "plugin cleanup completed")).toHaveLength(1); + expect(entries.filter((entry) => entry.message === "provider close completed")).toHaveLength(1); + await turn; + await client.archiveAgent(agent.id).catch(() => undefined); + } finally { + await writeFile(closeRelease, "release"); + await client.close(); + await daemon.close(); + await rm(cwd, { recursive: true, force: true }); + await rm(pluginDirectory, { recursive: true, force: true }); + } +}, 60_000); diff --git a/packages/server/src/server/plugins/plugin-process.ts b/packages/server/src/server/plugins/plugin-process.ts index 41310f78e8..1cc0f571a1 100644 --- a/packages/server/src/server/plugins/plugin-process.ts +++ b/packages/server/src/server/plugins/plugin-process.ts @@ -51,7 +51,7 @@ const handlers = new Map(); const providers = new Map(); const providerConnections = new Map< string, - { connection: ProviderConnection; unsubscribe: () => void } + { connection: ProviderConnection; unsubscribe: () => void; closing?: Promise } >(); const pendingProviderConnections = new Map(); let cleanup: (() => void | Promise) | null = null; @@ -193,6 +193,7 @@ async function sendProviderInput( if (stopping) throw new Error("Plugin is stopping"); const current = providerConnections.get(message.connectionId); if (!current) throw new Error(`Unknown provider connection: ${message.connectionId}`); + if (current.closing) throw new Error("Provider connection is closing"); await current.connection.send(message.input); send({ type: "provider.accepted", @@ -201,13 +202,25 @@ async function sendProviderInput( }); } +// The connection stays registered until its close has reported, so shutdown +// waits for a close already in flight instead of disconnecting underneath it. async function closeProviderConnection(connectionId: string): Promise { const current = providerConnections.get(connectionId); if (!current) return; - providerConnections.delete(connectionId); - current.unsubscribe(); - await current.connection.close(); - send({ type: "provider.closed", connectionId }); + if (current.closing) return current.closing; + const closing = (async () => { + current.unsubscribe(); + try { + await current.connection.close(); + send({ type: "provider.closed", connectionId }); + } catch (error) { + send({ type: "provider.closed", connectionId, error: describeError(error) }); + } finally { + providerConnections.delete(connectionId); + } + })(); + current.closing = closing; + return closing; } function runtimeRequire(name: string): unknown { @@ -397,13 +410,7 @@ process.on("message", (rawMessage: unknown) => { return; } if (message.type === "provider.close") { - void closeProviderConnection(message.connectionId).catch((error) => { - send({ - type: "provider.closed", - connectionId: message.connectionId, - error: describeError(error), - }); - }); + void closeProviderConnection(message.connectionId); return; } if (message.type === "paseo_frame" || message.type === "paseo_close") return; diff --git a/packages/server/src/server/plugins/runtime.posix.test.ts b/packages/server/src/server/plugins/runtime.posix.test.ts index 1f82d35909..c09ef9ebe2 100644 --- a/packages/server/src/server/plugins/runtime.posix.test.ts +++ b/packages/server/src/server/plugins/runtime.posix.test.ts @@ -34,6 +34,26 @@ async function readTextIfPresent(filePath: string): Promise { } } +function pluginMessages(runtime: PluginRuntime, pluginId: string): string[] { + return runtime.getLogs(pluginId).map((entry) => entry.message); +} + +function spawnObservedPluginProcess(onMessage: (message: unknown) => void) { + const loaderUrl = new URL("../../terminal/terminal-ts-loader.mjs", import.meta.url).href; + const setup = `import { register } from "node:module"; register(${JSON.stringify(loaderUrl)});`; + const child = fork(new URL("./plugin-process.ts", import.meta.url), [], { + execArgv: [ + "--experimental-strip-types", + "--import", + `data:text/javascript,${encodeURIComponent(setup)}`, + ], + serialization: "advanced", + stdio: ["ignore", "pipe", "pipe", "ipc"], + }); + child.on("message", onMessage); + return child; +} + function createReloadChild( name: string, events: string[], @@ -1040,6 +1060,75 @@ export default function contribute(plugin: unknown) { await rm(cleanupFile, { force: true }); }); + it("rejects provider input while a connection is closing", async () => { + const releaseDirectory = await mkdtemp(path.join(tmpdir(), "paseo-provider-close-")); + const releaseFile = path.join(releaseDirectory, "release"); + const directory = await createPlugin( + "closing-provider", + `import { readFile, writeFile } from "node:fs/promises"; +import type { ProviderRegistration } from "@getpaseo/plugin/server/provider"; +const provider: ProviderRegistration = { + id: "delayed", + label: "Delayed", + async connect() { + return { + version: 1, + capabilities: [], + async send() { console.log("provider accepted input"); }, + onEvent() { return () => undefined; }, + async close() { + console.log("provider close started"); + while (true) { + try { await readFile(${JSON.stringify(releaseFile)}); break; } + catch { await new Promise((resolve) => setTimeout(resolve, 10)); } + } + console.log("provider close completed"); + }, + }; + }, +}; +export default function contribute(server: { registerProvider(provider: ProviderRegistration): void }) { + server.registerProvider(provider); + return () => console.log("plugin cleanup started"); +}`, + ); + const closedReports: unknown[] = []; + const observeMessage = (message: unknown) => { + if ((message as { type?: string }).type === "provider.closed") closedReports.push(message); + }; + const runtime = createTestRuntime({ + spawnChild: () => spawnObservedPluginProcess(observeMessage), + }); + await runtime.startPlugin("closing-provider", directory); + try { + const connection = await runtime.connectProvider("closing-provider", "delayed", { + versions: [1], + capabilities: [], + }); + const closing = connection.close(); + await expect + .poll(() => pluginMessages(runtime, "closing-provider")) + .toContain("provider close started"); + await expect(connection.send({ type: "catalog", requestId: "during-close" })).rejects.toThrow( + "Provider connection is closing", + ); + const stopping = runtime.stopPluginById("closing-provider"); + await expect + .poll(() => pluginMessages(runtime, "closing-provider")) + .toContain("plugin cleanup started"); + await writeFile(releaseFile, "release"); + await Promise.all([closing, stopping]); + const messages = pluginMessages(runtime, "closing-provider"); + expect(messages).toContain("provider close completed"); + expect(messages).not.toContain("provider accepted input"); + expect(closedReports).toHaveLength(1); + } finally { + await writeFile(releaseFile, "release"); + await runtime.stopAll(); + await rm(releaseDirectory, { recursive: true, force: true }); + } + }); + it("closes a provider connection that resolves during plugin shutdown", async () => { const suffix = `${process.pid}-${Date.now()}`; const startedFile = path.join(tmpdir(), `paseo-provider-connect-started-${suffix}`);