diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 7635c58fc..c8d4e89c1 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -1,6 +1,6 @@ import { expect, test, vi } from "vitest"; import { spawn } from "node:child_process"; -import { mkdtempSync, rmSync } from "node:fs"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { join } from "node:path"; import { tmpdir } from "node:os"; import { randomUUID } from "node:crypto"; @@ -8,6 +8,7 @@ import { randomUUID } from "node:crypto"; import { createTestLogger } from "../../test-utils/test-logger.js"; import { AgentManager, + AgentManagerShuttingDownError, commandMayHaveChangedExternalState, type AgentManagerEvent, type ManagedAgent, @@ -139,6 +140,117 @@ class TestAgentClient implements AgentClient { } } +class HeldAgentCreationClient extends TestAgentClient { + private readonly creationStarted = deferred(); + private readonly creationAllowed = deferred(); + createdSessionClosed = false; + + override async createSession(config: AgentSessionConfig): Promise { + const recordSessionClosed = () => { + this.createdSessionClosed = true; + }; + const session = new (class extends TestAgentSession { + override async close(): Promise { + recordSessionClosed(); + } + })(config); + this.creationStarted.resolve(); + await this.creationAllowed.promise; + return session; + } + + waitForCreationToStart(): Promise { + return this.creationStarted.promise; + } + + finishCreating(): void { + this.creationAllowed.resolve(); + } +} + +class HeldAgentCreationAndCloseClient extends TestAgentClient { + private readonly creationStarted = deferred(); + private readonly creationAllowed = deferred(); + private readonly closeStarted = deferred(); + private readonly closeAllowed = deferred(); + + override async createSession(config: AgentSessionConfig): Promise { + this.creationStarted.resolve(); + await this.creationAllowed.promise; + const signalCloseStarted = () => this.closeStarted.resolve(); + const waitForClose = () => this.closeAllowed.promise; + return new (class extends TestAgentSession { + override async close(): Promise { + signalCloseStarted(); + await waitForClose(); + } + })(config); + } + + waitForCreationToStart(): Promise { + return this.creationStarted.promise; + } + + finishCreating(): void { + this.creationAllowed.resolve(); + } + + waitForCloseToStart(): Promise { + return this.closeStarted.promise; + } + + finishClosing(): void { + this.closeAllowed.resolve(); + } +} + +class HeldReloadCloseClient extends TestAgentClient { + private readonly closeStarted = deferred(); + private readonly closeAllowed = deferred(); + originalSessionClosed = false; + replacementSessionClosed = false; + + override async createSession(config: AgentSessionConfig): Promise { + const signalCloseStarted = () => this.closeStarted.resolve(); + const waitForClose = () => this.closeAllowed.promise; + const recordOriginalClosed = () => { + this.originalSessionClosed = true; + }; + return new (class extends TestAgentSession { + override async close(): Promise { + signalCloseStarted(); + await waitForClose(); + recordOriginalClosed(); + } + })(config); + } + + override async resumeSession( + _handle: AgentPersistenceHandle, + config?: Partial, + ): Promise { + const recordReplacementClosed = () => { + this.replacementSessionClosed = true; + }; + return new (class extends TestAgentSession { + override async close(): Promise { + recordReplacementClosed(); + } + })({ + provider: "codex", + cwd: config?.cwd ?? process.cwd(), + }); + } + + waitForCloseToStart(): Promise { + return this.closeStarted.promise; + } + + finishClosing(): void { + this.closeAllowed.resolve(); + } +} + class NativeArchiveRecordingClient extends TestAgentClient { readonly archivedHandles: AgentPersistenceHandle[] = []; readonly unarchivedHandles: AgentPersistenceHandle[] = []; @@ -290,6 +402,52 @@ class TestAgentSession implements AgentSession { async close(): Promise {} } +class HeldRuntimeInfoSession extends TestAgentSession { + private readonly runtimeInfoRequested = deferred(); + private readonly runtimeInfoAllowed = deferred(); + + override async getRuntimeInfo() { + this.runtimeInfoRequested.resolve(); + await this.runtimeInfoAllowed.promise; + return await super.getRuntimeInfo(); + } + + waitForRuntimeInfo(): Promise { + return this.runtimeInfoRequested.promise; + } + + finishRuntimeInfo(): void { + this.runtimeInfoAllowed.resolve(); + } +} + +class HeldRuntimeInfoClient extends TestAgentClient { + private readonly sessionCreated = deferred(); + private session: HeldRuntimeInfoSession | null = null; + + override async createSession(config: AgentSessionConfig): Promise { + this.session = new HeldRuntimeInfoSession(config); + this.sessionCreated.resolve(this.session); + return this.session; + } + + async waitForRuntimeInfo(): Promise { + const session = await this.sessionCreated.promise; + await session.waitForRuntimeInfo(); + } + + finishRuntimeInfo(): void { + this.requireSession().finishRuntimeInfo(); + } + + private requireSession(): HeldRuntimeInfoSession { + if (!this.session) { + throw new Error("Expected a created session"); + } + return this.session; + } +} + class StreamingAssistantSession implements AgentSession { readonly provider = "codex" as const; readonly capabilities = TEST_CAPABILITIES; @@ -449,6 +607,211 @@ function fakeCodexEmitting(args: FakeCodexEmitterArgs): AgentClient { const logger = createTestLogger(); +test("does not register a session that finishes starting after shutdown begins", async () => { + const client = new HeldAgentCreationClient(); + const manager = new AgentManager({ + clients: { codex: client }, + logger, + idFactory: () => "00000000-0000-4000-8000-000000000100", + }); + + const creation = manager.createAgent( + { + provider: "codex", + cwd: process.cwd(), + }, + undefined, + { workspaceId: undefined }, + ); + await client.waitForCreationToStart(); + + manager.prepareForShutdown(); + client.finishCreating(); + + await expect(creation).rejects.toThrow("Agent manager is shutting down"); + expect({ agents: manager.listAgents(), sessionClosed: client.createdSessionClosed }).toEqual({ + agents: [], + sessionClosed: true, + }); +}); + +test("flush waits for rejected session cleanup that starts after shutdown", async () => { + const client = new HeldAgentCreationAndCloseClient(); + const manager = new AgentManager({ + clients: { codex: client }, + logger, + idFactory: () => "00000000-0000-4000-8000-000000000098", + }); + + const creation = manager + .createAgent( + { + provider: "codex", + cwd: process.cwd(), + }, + undefined, + { workspaceId: undefined }, + ) + .catch((error: unknown) => error); + await client.waitForCreationToStart(); + + manager.prepareForShutdown(); + let flushResolved = false; + const flushing = manager.flushForShutdown().then(() => { + flushResolved = true; + return undefined; + }); + client.finishCreating(); + await client.waitForCloseToStart(); + + try { + expect(flushResolved).toBe(false); + } finally { + client.finishClosing(); + } + + expect(await creation).toBeInstanceOf(AgentManagerShuttingDownError); + await flushing; + expect(manager.listAgents()).toEqual([]); +}); + +test("does not persist an initializing session after shutdown closes it", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-shutdown-register-test-")); + const storage = new AgentStorage(join(workdir, "agents"), logger); + const client = new HeldRuntimeInfoClient(); + const agentId = "00000000-0000-4000-8000-000000000099"; + const manager = new AgentManager({ + clients: { codex: client }, + registry: storage, + logger, + idFactory: () => agentId, + }); + + try { + const creation = manager.createAgent( + { + provider: "codex", + cwd: workdir, + }, + undefined, + { workspaceId: undefined }, + ); + await client.waitForRuntimeInfo(); + + manager.prepareForShutdown(); + const closing = manager.closeAgent(agentId); + client.finishRuntimeInfo(); + + await expect(creation).rejects.toBeInstanceOf(AgentManagerShuttingDownError); + await closing; + await storage.flush(); + expect({ agents: manager.listAgents(), record: await storage.get(agentId) }).toMatchObject({ + agents: [], + record: { lastStatus: "closed" }, + }); + } finally { + rmSync(workdir, { recursive: true, force: true }); + } +}); + +test("reload leaves a closed durable snapshot when shutdown starts during the swap", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-shutdown-reload-test-")); + const storage = new AgentStorage(join(workdir, "agents"), logger); + const client = new HeldReloadCloseClient(); + const agentId = "00000000-0000-4000-8000-000000000097"; + const manager = new AgentManager({ + clients: { codex: client }, + registry: storage, + logger, + idFactory: () => agentId, + }); + + try { + await manager.createAgent( + { + provider: "codex", + cwd: workdir, + }, + undefined, + { workspaceId: undefined }, + ); + const reload = manager.reloadAgentSession(agentId).catch((error: unknown) => error); + await client.waitForCloseToStart(); + + manager.prepareForShutdown(); + const closing = Promise.all(manager.listAgents().map((agent) => manager.closeAgent(agent.id))); + client.finishClosing(); + + await closing; + expect(await reload).toBeInstanceOf(AgentManagerShuttingDownError); + await manager.flush(); + await storage.flush(); + expect({ + agents: manager.listAgents(), + record: await storage.get(agentId), + replacementSessionClosed: client.replacementSessionClosed, + }).toMatchObject({ + agents: [], + record: { lastStatus: "closed" }, + replacementSessionClosed: true, + }); + } finally { + client.finishClosing(); + await manager.flush().catch(() => undefined); + await storage.flush().catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + +test("reload closes both sessions when the closed snapshot cannot be persisted", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-reload-persist-failure-test-")); + const storagePath = join(workdir, "agents"); + const storage = new AgentStorage(storagePath, logger); + const client = new HeldReloadCloseClient(); + const agentId = "00000000-0000-4000-8000-000000000096"; + const manager = new AgentManager({ + clients: { codex: client }, + registry: storage, + logger, + idFactory: () => agentId, + }); + + try { + await manager.createAgent( + { + provider: "codex", + cwd: workdir, + }, + undefined, + { workspaceId: undefined }, + ); + await storage.flush(); + rmSync(storagePath, { recursive: true, force: true }); + writeFileSync(storagePath, "blocks the storage directory"); + + const reload = manager.reloadAgentSession(agentId).catch((error: unknown) => error); + await client.waitForCloseToStart(); + client.finishClosing(); + + expect(await reload).toBeInstanceOf(Error); + await manager.flushForShutdown(); + expect({ + agents: manager.listAgents(), + originalSessionClosed: client.originalSessionClosed, + replacementSessionClosed: client.replacementSessionClosed, + }).toEqual({ + agents: [], + originalSessionClosed: true, + replacementSessionClosed: true, + }); + } finally { + client.finishClosing(); + await manager.flushForShutdown().catch(() => undefined); + await storage.flush().catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + test("normalizeConfig injects the provider default model when omitted", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-test-")); const storagePath = join(workdir, "agents"); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 48ad8bcf7..232f5c57a 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -82,6 +82,13 @@ const STORED_AGENT_CAPABILITIES: AgentCapabilityFlags = { type TimeoutResult = "completed" | "timed_out"; +export class AgentManagerShuttingDownError extends Error { + constructor() { + super("Agent manager is shutting down"); + this.name = "AgentManagerShuttingDownError"; + } +} + interface PreparedSessionConfig { storedConfig: AgentSessionConfig; launchConfig: AgentSessionConfig; @@ -535,6 +542,7 @@ export class AgentManager { private readonly durableTimelineStore?: AgentTimelineStore; private readonly previousStatuses = new Map(); private readonly backgroundTasks = new Set>(); + private readonly agentRegistrationTasks = new Set>(); private readonly agentStreamCoalescer: AgentStreamCoalescer; private mcpBaseUrl: string | null; private readonly mcpAuthToken: string | null; @@ -546,6 +554,7 @@ export class AgentManager { private onWorkspaceStateMayHaveChanged?: (params: { cwd: string }) => void; private logger: Logger; private readonly rescueTimeouts: Required; + private acceptingAgentRegistrations = true; constructor(options: AgentManagerOptions) { this.idFactory = options?.idFactory ?? (() => randomUUID()); @@ -619,6 +628,10 @@ export class AgentManager { this.mcpBaseUrl = url; } + prepareForShutdown(): void { + this.acceptingAgentRegistrations = false; + } + setPaseoToolsEnabled(enabled: boolean): void { this.paseoToolsEnabled = enabled; } @@ -922,11 +935,20 @@ export class AgentManager { return this.timelineStore.fetch(id, options); } - async createAgent( + createAgent( config: AgentSessionConfig, agentId: string | undefined, options: CreateAgentOptions, ): Promise { + return this.trackAgentRegistrationOperation(this.createAgentInternal(config, agentId, options)); + } + + private async createAgentInternal( + config: AgentSessionConfig, + agentId: string | undefined, + options: CreateAgentOptions, + ): Promise { + this.assertAcceptingAgentRegistrations(); const resolvedAgentId = validateAgentId(agentId ?? this.idFactory(), "createAgent"); const { storedConfig, launchConfig } = await this.prepareSessionConfig(config, resolvedAgentId); this.requireEnabledProvider(storedConfig.provider); @@ -954,7 +976,7 @@ export class AgentManager { // Reconstruct an agent from provider persistence. Callers should explicitly // hydrate timeline history after resume. - async resumeAgentFromPersistence( + resumeAgentFromPersistence( handle: AgentPersistenceHandle, overrides?: Partial, agentId?: string, @@ -966,6 +988,24 @@ export class AgentManager { workspaceId?: string; }, ): Promise { + return this.trackAgentRegistrationOperation( + this.resumeAgentFromPersistenceInternal(handle, overrides, agentId, options), + ); + } + + private async resumeAgentFromPersistenceInternal( + handle: AgentPersistenceHandle, + overrides?: Partial, + agentId?: string, + options?: { + createdAt?: Date; + updatedAt?: Date; + lastUserMessageAt?: Date | null; + labels?: Record; + workspaceId?: string; + }, + ): Promise { + this.assertAcceptingAgentRegistrations(); const resolvedAgentId = validateAgentId( agentId ?? this.idFactory(), "resumeAgentFromPersistence", @@ -994,13 +1034,24 @@ export class AgentManager { return this.registerSession(session, storedConfig, resolvedAgentId, options); } - async importProviderSession(input: { + importProviderSession(input: { provider: AgentProvider; providerHandleId: string; cwd: string; workspaceId: string; labels?: Record; }): Promise { + return this.trackAgentRegistrationOperation(this.importProviderSessionInternal(input)); + } + + private async importProviderSessionInternal(input: { + provider: AgentProvider; + providerHandleId: string; + cwd: string; + workspaceId: string; + labels?: Record; + }): Promise { + this.assertAcceptingAgentRegistrations(); const resolvedAgentId = validateAgentId(this.idFactory(), "importProviderSession"); this.requireEnabledProvider(input.provider); @@ -1025,20 +1076,30 @@ export class AgentManager { }, { config: providerLaunchConfig, storedConfig, launchContext }, ); - const importedConfig = await this.normalizeConfig(stripInternalPaseoMcpServer(imported.config)); - const timelineRows = buildImportedTimelineRows(imported.timeline); - const initialTitle = resolveImportedAgentTitle(importedConfig, timelineRows); + let handedToRegistration = false; + try { + const importedConfig = await this.normalizeConfig( + stripInternalPaseoMcpServer(imported.config), + ); + const timelineRows = buildImportedTimelineRows(imported.timeline); + const initialTitle = resolveImportedAgentTitle(importedConfig, timelineRows); - return this.registerSession(imported.session, importedConfig, resolvedAgentId, { - labels: input.labels, - workspaceId: input.workspaceId, - timelineRows, - timelineNextSeq: timelineRows.length + 1, - persistence: imported.persistence, - historyPrimed: true, - initialTitle, - publishWhenReady: true, - }); + handedToRegistration = true; + return this.registerSession(imported.session, importedConfig, resolvedAgentId, { + labels: input.labels, + workspaceId: input.workspaceId, + timelineRows, + timelineNextSeq: timelineRows.length + 1, + persistence: imported.persistence, + historyPrimed: true, + initialTitle, + publishWhenReady: true, + }); + } finally { + if (!handedToRegistration) { + await this.closeUnregisteredSession(imported.session); + } + } } // Hot-reload an active agent session with config overrides. By default the @@ -1047,11 +1108,22 @@ export class AgentManager { // new epoch is minted and provider history is re-streamed — this is what the // user-facing "Reload agent" action wants when the on-disk session was // mutated outside Paseo. - async reloadAgentSession( + reloadAgentSession( agentId: string, overrides?: Partial, options?: { rehydrateFromDisk?: boolean }, ): Promise { + return this.trackAgentRegistrationOperation( + this.reloadAgentSessionInternal(agentId, overrides, options), + ); + } + + private async reloadAgentSessionInternal( + agentId: string, + overrides?: Partial, + options?: { rehydrateFromDisk?: boolean }, + ): Promise { + this.assertAcceptingAgentRegistrations(); let existing = this.requireSessionAgent(agentId); if (this.hasInFlightRun(agentId)) { await this.cancelAgentRun(agentId); @@ -1078,36 +1150,43 @@ export class AgentManager { ? await client.resumeSession(handle, providerLaunchConfig, launchContext) : await client.createSession(providerLaunchConfig, launchContext); - this.agentStreamCoalescer.flushAndDiscard(agentId); - // Remove the existing agent entry before swapping sessions - this.agents.delete(agentId); - if (existing.unsubscribeSession) { - existing.unsubscribeSession(); - existing.unsubscribeSession = null; - } - this.foregroundRuns.clearAgent(agentId, existing); - await this.closeReloadedSession(existing.session, agentId); + let handedToRegistration = false; + try { + this.assertAcceptingAgentRegistrations(); - if (rehydrateFromDisk) { - // Wipe both durable and in-memory timeline so registerSession mints a - // new epoch and hydrateTimelineFromProvider re-streams the freshly read - // provider history into an empty timeline. - await this.deleteCommittedTimeline(agentId); - this.timelineStore.delete(agentId); - } + const closedExisting = this.prepareAgentForClosure(existing, "agent reloaded"); + try { + await this.persistSnapshot(closedExisting); + } finally { + await this.closeReloadedSession(existing.session, agentId); + } - // Preserve existing labels and timeline during reload. - return this.registerSession(session, storedConfig, agentId, { - labels: existing.labels, - workspaceId: existing.workspaceId, - createdAt: existing.createdAt, - updatedAt: existing.updatedAt, - lastUserMessageAt: existing.lastUserMessageAt, - historyPrimed: rehydrateFromDisk ? false : preservedHistoryPrimed, - lastUsage: preservedLastUsage, - lastError: preservedLastError, - attention: preservedAttention, - }); + if (rehydrateFromDisk) { + // Wipe both durable and in-memory timeline so registerSession mints a + // new epoch and hydrateTimelineFromProvider re-streams the freshly read + // provider history into an empty timeline. + await this.deleteCommittedTimeline(agentId); + this.timelineStore.delete(agentId); + } + + // Preserve existing labels and timeline during reload. + handedToRegistration = true; + return this.registerSession(session, storedConfig, agentId, { + labels: existing.labels, + workspaceId: existing.workspaceId, + createdAt: existing.createdAt, + updatedAt: existing.updatedAt, + lastUserMessageAt: existing.lastUserMessageAt, + historyPrimed: rehydrateFromDisk ? false : preservedHistoryPrimed, + lastUsage: preservedLastUsage, + lastError: preservedLastError, + attention: preservedAttention, + }); + } finally { + if (!handedToRegistration) { + await this.closeUnregisteredSession(session); + } + } } private async closeReloadedSession(session: AgentSession, agentId: string): Promise { @@ -2474,49 +2553,84 @@ export class AgentManager { workspaceId?: string; }, ): Promise { - const resolvedAgentId = validateAgentId(agentId, "registerSession"); - if (this.agents.has(resolvedAgentId)) { - throw new Error(`Agent with id ${resolvedAgentId} already exists`); - } - const initialPersistedTitle = await this.resolveInitialPersistedTitle( - resolvedAgentId, - config, - options?.initialTitle ?? null, - ); + let registered = false; + try { + this.assertAcceptingAgentRegistrations(); + const resolvedAgentId = validateAgentId(agentId, "registerSession"); + if (this.agents.has(resolvedAgentId)) { + throw new Error(`Agent with id ${resolvedAgentId} already exists`); + } + const initialPersistedTitle = await this.resolveInitialPersistedTitle( + resolvedAgentId, + config, + options?.initialTitle ?? null, + ); - const now = new Date(); - const { durableTimelineHasRows } = await this.initializeAgentTimelineForRegister({ - agentId: resolvedAgentId, - now, - options, - }); + const now = new Date(); + const { durableTimelineHasRows } = await this.initializeAgentTimelineForRegister({ + agentId: resolvedAgentId, + now, + options, + }); - const managed = this.buildManagedAgentForRegister({ - resolvedAgentId, - session, - config, - now, - durableTimelineHasRows, - options, - }); + const managed = this.buildManagedAgentForRegister({ + resolvedAgentId, + session, + config, + now, + durableTimelineHasRows, + options, + }); - this.agents.set(resolvedAgentId, managed); - // Initialize previousStatus to track transitions - this.previousStatuses.set(resolvedAgentId, managed.lifecycle); - await this.refreshRuntimeInfo(managed, { emit: !options?.publishWhenReady }); - await this.persistSnapshot(managed, { - title: initialPersistedTitle, - }); - if (!options?.publishWhenReady) { + this.assertAcceptingAgentRegistrations(); + this.agents.set(resolvedAgentId, managed); + registered = true; + // Initialize previousStatus to track transitions + this.previousStatuses.set(resolvedAgentId, managed.lifecycle); + await this.refreshRuntimeInfo(managed, { emit: false }); + this.assertAgentRegistrationActive(managed); + await this.persistSnapshot(managed, { + title: initialPersistedTitle, + }); + this.assertAgentRegistrationActive(managed); + if (!options?.publishWhenReady) { + this.emitState(managed, { persist: false }); + } + + await this.refreshSessionState(managed, { emit: false }); + this.assertAgentRegistrationActive(managed); + managed.lifecycle = "idle"; + await this.persistSnapshot(managed); + this.assertAgentRegistrationActive(managed); this.emitState(managed, { persist: false }); + this.subscribeToSession(managed); + return { ...managed }; + } catch (error) { + if (!registered) { + await this.closeUnregisteredSession(session); + } + throw error; } + } - await this.refreshSessionState(managed, { emit: !options?.publishWhenReady }); - managed.lifecycle = "idle"; - await this.persistSnapshot(managed); - this.emitState(managed, { persist: false }); - this.subscribeToSession(managed); - return { ...managed }; + private assertAcceptingAgentRegistrations(): void { + if (!this.acceptingAgentRegistrations) { + throw new AgentManagerShuttingDownError(); + } + } + + private assertAgentRegistrationActive(agent: ActiveManagedAgent): void { + if (!this.acceptingAgentRegistrations || this.agents.get(agent.id) !== agent) { + throw new AgentManagerShuttingDownError(); + } + } + + private async closeUnregisteredSession(session: AgentSession): Promise { + try { + await session.close(); + } catch (error) { + this.logger.warn({ err: error }, "Failed to close unregistered agent session"); + } } private async initializeAgentTimelineForRegister(params: { @@ -3565,15 +3679,45 @@ export class AgentManager { }); } + private trackAgentRegistrationOperation(result: Promise): Promise { + const settled = result.then( + () => undefined, + () => undefined, + ); + this.agentRegistrationTasks.add(settled); + void settled.then(() => { + this.agentRegistrationTasks.delete(settled); + return undefined; + }); + return result; + } + /** * Flush any background persistence work (best-effort). - * Used by daemon shutdown paths to avoid unhandled rejections after cleanup. */ async flush(): Promise { + await this.flushTasks({ includeAgentRegistrations: false }); + } + + /** + * Flush persistence and agent registrations that crossed the synchronous + * shutdown barrier. Those registrations own provider sessions until they + * either install them or close them. + */ + async flushForShutdown(): Promise { + await this.flushTasks({ includeAgentRegistrations: true }); + } + + private async flushTasks(options: { includeAgentRegistrations: boolean }): Promise { this.agentStreamCoalescer.flushAll(); // Drain tasks, including tasks spawned while awaiting. - while (this.backgroundTasks.size > 0) { - const pending = Array.from(this.backgroundTasks); + while ( + this.backgroundTasks.size > 0 || + (options.includeAgentRegistrations && this.agentRegistrationTasks.size > 0) + ) { + const pending = options.includeAgentRegistrations + ? [...this.backgroundTasks, ...this.agentRegistrationTasks] + : [...this.backgroundTasks]; await Promise.allSettled(pending); } } diff --git a/packages/server/src/server/bootstrap.smoke.test.ts b/packages/server/src/server/bootstrap.smoke.test.ts index d4e6bbff5..5f762fa24 100644 --- a/packages/server/src/server/bootstrap.smoke.test.ts +++ b/packages/server/src/server/bootstrap.smoke.test.ts @@ -7,6 +7,7 @@ import { afterEach, describe, expect, test, vi } from "vitest"; import { WebSocket } from "ws"; import { createPaseoDaemon, parseListenString, type PaseoDaemonConfig } from "./bootstrap.js"; +import { AgentManagerShuttingDownError } from "./agent/agent-manager.js"; import { hashDaemonPassword } from "./auth.js"; import { generateLocalPairingOffer } from "./pairing-offer.js"; import { createTestPaseoDaemon } from "./test-utils/paseo-daemon.js"; @@ -14,6 +15,23 @@ import { createTestAgentClients } from "./test-utils/fake-agent-client.js"; import { isPlatform } from "../test-utils/platform.js"; import { findFreePort } from "./service-proxy.js"; +interface HeldAgentClose { + started: Promise; + arm(): void; + closeSession(): Promise; + finish(): void; +} + +interface BlockedDaemonShutdown { + probeReconnect(): Promise; + tryCreateAgent(): Promise<"created" | "rejected">; + finish(): Promise; +} + +type WebSocketProbeResult = + | { status: "connected" } + | { status: "rejected"; statusCode: number | null }; + describe("paseo daemon bootstrap", () => { afterEach(() => { vi.restoreAllMocks(); @@ -206,6 +224,17 @@ describe("paseo daemon bootstrap", () => { } }); + test("stops new connections and agent registrations before closing agents", async () => { + const shutdown = await beginDaemonShutdownWithAgentClosing(); + try { + await expect( + Promise.all([shutdown.probeReconnect(), shutdown.tryCreateAgent()]), + ).resolves.toEqual([{ status: "rejected", statusCode: 503 }, "rejected"]); + } finally { + await shutdown.finish(); + } + }); + test("standalone listener exposes services only", async () => { const standalonePort = await findFreePort(); const upstream = http.createServer((_req, res) => { @@ -522,3 +551,97 @@ describe("paseo daemon bootstrap", () => { }, ); }); + +function holdAgentClose(): HeldAgentClose { + let armed = false; + let markStarted = () => {}; + let finish = () => {}; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + const finished = new Promise((resolve) => { + finish = resolve; + }); + return { + started, + arm() { + armed = true; + }, + async closeSession() { + if (!armed) { + return; + } + markStarted(); + await finished; + }, + finish: () => finish(), + }; +} + +async function beginDaemonShutdownWithAgentClosing(): Promise { + const heldAgentClose = holdAgentClose(); + const daemonHandle = await createTestPaseoDaemon({ + cleanup: false, + agentClients: createTestAgentClients({ closeSession: heldAgentClose.closeSession }), + }); + const agentCwd = await mkdtemp(path.join(os.tmpdir(), "paseo-shutdown-agent-")); + await daemonHandle.daemon.agentManager.createAgent( + { + provider: "codex", + cwd: agentCwd, + }, + undefined, + { workspaceId: undefined }, + ); + + heldAgentClose.arm(); + const stopPromise = daemonHandle.daemon.stop(); + await heldAgentClose.started; + + return { + probeReconnect: () => probeWebSocketConnection(`ws://127.0.0.1:${daemonHandle.port}/ws`), + async tryCreateAgent() { + try { + await daemonHandle.daemon.agentManager.createAgent( + { + provider: "codex", + cwd: agentCwd, + }, + undefined, + { workspaceId: undefined }, + ); + return "created"; + } catch (error) { + if (error instanceof AgentManagerShuttingDownError) { + return "rejected"; + } + throw error; + } + }, + async finish() { + heldAgentClose.finish(); + await stopPromise; + await daemonHandle.daemon.agentManager.flush().catch(() => undefined); + await Promise.all([ + rm(path.dirname(daemonHandle.paseoHome), { recursive: true, force: true }), + rm(daemonHandle.staticDir, { recursive: true, force: true }), + rm(agentCwd, { recursive: true, force: true }), + ]); + }, + }; +} + +function probeWebSocketConnection(url: string): Promise { + const ws = new WebSocket(url); + return new Promise((resolve) => { + ws.once("open", () => { + ws.close(); + resolve({ status: "connected" }); + }); + ws.once("error", () => resolve({ status: "rejected", statusCode: null })); + ws.once("unexpected-response", (_request, response) => { + response.resume(); + resolve({ status: "rejected", statusCode: response.statusCode ?? null }); + }); + }); +} diff --git a/packages/server/src/server/bootstrap.ts b/packages/server/src/server/bootstrap.ts index 0d86b08f7..2c0b2885f 100644 --- a/packages/server/src/server/bootstrap.ts +++ b/packages/server/src/server/bootstrap.ts @@ -1443,8 +1443,11 @@ export async function createPaseoDaemon( const stop = async () => { scriptHealthMonitor.stop(); + // Freeze both ingress and registration before taking the agent closure snapshot. + wsServer?.prepareForShutdown(); + agentManager.prepareForShutdown(); await closeAllAgents(logger, agentManager); - await agentManager.flush().catch(() => undefined); + await agentManager.flushForShutdown().catch(() => undefined); detachAgentStoragePersistence(); await agentStorage.flush().catch(() => undefined); await providerSnapshotManager.shutdown(); diff --git a/packages/server/src/server/test-utils/fake-agent-client.ts b/packages/server/src/server/test-utils/fake-agent-client.ts index 81335f6f8..6468f9fb1 100644 --- a/packages/server/src/server/test-utils/fake-agent-client.ts +++ b/packages/server/src/server/test-utils/fake-agent-client.ts @@ -51,6 +51,18 @@ interface Deferred { reject: (err: unknown) => void; } +interface FakeAgentSessionOptions { + providerName: string; + config: AgentSessionConfig; + sessionId?: string; + memoryMarker?: string | null; + closeSession?: () => Promise; +} + +export interface TestAgentClientOptions { + closeSession?: () => Promise; +} + function createDeferred(): Deferred { let resolve!: (value: T) => void; let reject!: (err: unknown) => void; @@ -316,16 +328,14 @@ class FakeAgentSession implements AgentSession { private nextTurnOrdinal = 0; private activeForegroundTurnId: string | null = null; - constructor( - providerName: string, - config: AgentSessionConfig, - sessionId?: string, - memoryMarker?: string | null, - ) { - this.providerName = providerName; - this.config = config; - this.id = sessionId ?? randomUUID(); - this.memoryMarker = memoryMarker ?? null; + private readonly closeSession: (() => Promise) | undefined; + + constructor(options: FakeAgentSessionOptions) { + this.providerName = options.providerName; + this.config = options.config; + this.id = options.sessionId ?? randomUUID(); + this.memoryMarker = options.memoryMarker ?? null; + this.closeSession = options.closeSession; this.historyPath = path.join( tmpdir(), "paseo-fake-provider-history", @@ -847,7 +857,9 @@ class FakeAgentSession implements AgentSession { this.interruptSignal.resolve(); } - async close(): Promise {} + async close(): Promise { + await this.closeSession?.(); + } async listCommands(): Promise { if (this.providerName === "codex") { @@ -1155,13 +1167,20 @@ class FakeAgentSession implements AgentSession { class FakeAgentClient implements AgentClient { readonly capabilities = TEST_CAPABILITIES; - constructor(public readonly provider: string) {} + constructor( + public readonly provider: string, + private readonly options: TestAgentClientOptions, + ) {} async createSession( config: AgentSessionConfig, _launchContext?: AgentLaunchContext, ): Promise { - return new FakeAgentSession(this.provider, { ...config }); + return new FakeAgentSession({ + providerName: this.provider, + config: { ...config }, + closeSession: this.options.closeSession, + }); } async resumeSession( @@ -1178,12 +1197,13 @@ class FakeAgentClient implements AgentClient { (handle.metadata as Record | undefined)?.marker ?? (handle.metadata as Record | undefined)?.conversationId ?? null; - return new FakeAgentSession( - this.provider, - cfg, - handle.sessionId, - typeof marker === "string" ? marker : null, - ); + return new FakeAgentSession({ + providerName: this.provider, + config: cfg, + sessionId: handle.sessionId, + memoryMarker: typeof marker === "string" ? marker : null, + closeSession: this.options.closeSession, + }); } async fetchCatalog( @@ -1222,10 +1242,12 @@ class FakeAgentClient implements AgentClient { } } -export function createTestAgentClients(): Record { +export function createTestAgentClients( + options: TestAgentClientOptions = {}, +): Record { return { - claude: new FakeAgentClient("claude"), - codex: new FakeAgentClient("codex"), - opencode: new FakeAgentClient("opencode"), + claude: new FakeAgentClient("claude", options), + codex: new FakeAgentClient("codex", options), + opencode: new FakeAgentClient("opencode", options), }; } diff --git a/packages/server/src/server/websocket-server.relay-reconnect.test.ts b/packages/server/src/server/websocket-server.relay-reconnect.test.ts index 64e2918fe..be49b8dae 100644 --- a/packages/server/src/server/websocket-server.relay-reconnect.test.ts +++ b/packages/server/src/server/websocket-server.relay-reconnect.test.ts @@ -454,6 +454,21 @@ function holdNextSessionMessage(session: (typeof sessionMock.instances)[number]) }; } +function holdSessionCleanup(session: (typeof sessionMock.instances)[number]): { + finish: () => void; +} { + let finish = () => {}; + session.cleanup.mockImplementationOnce( + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + return { + finish: () => finish(), + }; +} + describe("relay external socket reconnect behavior", () => { beforeEach(() => { sessionMock.instances.length = 0; @@ -517,6 +532,36 @@ describe("relay external socket reconnect behavior", () => { await server.close(); }); + test("rejects sockets attached after shutdown begins", async () => { + const server = createServer(); + const existingSocket = new MockSocket(); + await attachRelayAndHello({ + server, + socket: existingSocket, + clientId: "existing-client", + }); + + const heldCleanup = holdSessionCleanup(sessionMock.instances[0]); + const closePromise = server.close(); + + const lateSocket = new MockSocket(); + try { + await server.attachExternalSocket(lateSocket, { transport: "relay" }); + lateSocket.emit("message", JSON.stringify(createHelloMessage("late-client"))); + + expect({ + readyState: lateSocket.readyState, + sessionCount: sessionMock.instances.length, + }).toEqual({ + readyState: 3, + sessionCount: 1, + }); + } finally { + heldCleanup.finish(); + await closePromise; + } + }); + test("closes pending connection when hello timeout elapses", async () => { const server = createServer(); diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 3c2143b9b..5b45cf7a3 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -351,6 +351,7 @@ const HELLO_TIMEOUT_MS = 15_000; const WS_CLOSE_HELLO_TIMEOUT = 4001; const WS_CLOSE_INVALID_HELLO = 4002; const WS_CLOSE_INCOMPATIBLE_PROTOCOL = 4003; +const WS_CLOSE_SERVER_SHUTDOWN = 1001; const WS_PROTOCOL_VERSION = 1; const WS_RUNTIME_METRICS_FLUSH_MS = 30_000; @@ -463,6 +464,7 @@ export class VoiceAssistantWebSocketServer { private unsubscribeTerminalActivity: (() => void) | null = null; private readonly browserToolsBroker: BrowserToolsBroker | null; private readonly browserToolsRegistrations = new Map(); + private acceptingConnections = true; constructor( server: HTTPServer, @@ -708,6 +710,11 @@ export class VoiceAssistantWebSocketServer { hostnames: HostnamesConfig | undefined, callback: (res: boolean, code?: number, message?: string) => void, ): void { + if (!this.acceptingConnections) { + callback(false, 503, "Server shutting down"); + return; + } + const requestMetadata = extractSocketRequestMetadata(req); const origin = requestMetadata.origin; const requestHost = requestMetadata.host ?? null; @@ -799,7 +806,12 @@ export class VoiceAssistantWebSocketServer { await this.attachSocket(ws, undefined, metadata); } + public prepareForShutdown(): void { + this.acceptingConnections = false; + } + public async close(): Promise { + this.prepareForShutdown(); this.unsubscribeSpeechReadiness?.(); this.unsubscribeSpeechReadiness = null; this.unsubscribeDaemonConfigChange?.(); @@ -921,6 +933,15 @@ export class VoiceAssistantWebSocketServer { request?: unknown, metadata?: ExternalSocketMetadata, ): Promise { + if (!this.acceptingConnections) { + try { + ws.close(WS_CLOSE_SERVER_SHUTDOWN, "Server shutting down"); + } catch { + // ignore close errors + } + return; + } + const requestMetadata = extractSocketRequestMetadata(request); const identity = createWebSocketConnectionIdentity(requestMetadata, metadata); this.socketIdentities.set(ws, identity); @@ -1609,6 +1630,10 @@ export class VoiceAssistantWebSocketServer { ws: WebSocketLike, data: Buffer | ArrayBuffer | Buffer[] | string, ): void { + if (!this.acceptingConnections) { + return; + } + const activeConnection = this.sessions.get(ws); const pendingConnection = this.pendingConnections.get(ws); const log =