diff --git a/apps/host-daemon/src/runtime-manager.test.ts b/apps/host-daemon/src/runtime-manager.test.ts index ee4518dfb5b..8c37ba6a848 100644 --- a/apps/host-daemon/src/runtime-manager.test.ts +++ b/apps/host-daemon/src/runtime-manager.test.ts @@ -54,10 +54,7 @@ interface RuntimeOptionsRef { } interface RuntimeManagerProviderMaintenanceInternals { - createProviderMaintenanceRuntime: (args: { - dataDir: string; - }) => Promise; - providerMaintenanceRuntime: AgentRuntime | null; + createProviderMaintenanceRuntime: () => Promise; } const execFileAsync = promisify(execFile); @@ -1161,6 +1158,104 @@ describe("RuntimeManager", () => { } }); + it("retires a busy provider maintenance runtime after its requests finish when the shell env changes", async () => { + const dataDir = await makeTempDir("bb-provider-maintenance-retire-"); + const firstRuntime = createFakeRuntime(); + const secondRuntime = createFakeRuntime(); + const createRuntime = vi + .fn() + .mockReturnValueOnce(firstRuntime) + .mockReturnValueOnce(secondRuntime); + const manager = new RuntimeManager({ + createRuntime, + shellEnv: { PATH: "/old/bin:/usr/bin" }, + }); + const usage = createDeferredPromise(); + const usageStarted = createDeferredPromise(); + + const usageRequest = manager.withProviderMaintenanceRuntime( + { dataDir }, + async (runtime) => { + usageStarted.resolve(runtime); + return usage.promise; + }, + ); + await expect(usageStarted.promise).resolves.toBe(firstRuntime); + + await manager.replaceBaseShellEnv({ PATH: "/new/bin:/usr/bin" }); + expect(firstRuntime.shutdown).not.toHaveBeenCalled(); + await expect( + manager.withProviderMaintenanceRuntime({ dataDir }, async (runtime) => + runtime === secondRuntime ? "fresh" : "stale", + ), + ).resolves.toBe("fresh"); + + usage.resolve("usage"); + await expect(usageRequest).resolves.toBe("usage"); + await vi.waitFor(() => { + expect(firstRuntime.shutdown).toHaveBeenCalledOnce(); + }); + expect(secondRuntime.shutdown).not.toHaveBeenCalled(); + + await manager.shutdownAll(); + expect(secondRuntime.shutdown).toHaveBeenCalledOnce(); + }); + + it("shuts down a busy retired provider maintenance runtime when the manager shuts down", async () => { + const dataDir = await makeTempDir("bb-provider-maintenance-shutdown-"); + const runtime = createFakeRuntime(); + const manager = new RuntimeManager({ + createRuntime: () => runtime, + shellEnv: { PATH: "/old/bin:/usr/bin" }, + }); + const requestStarted = createDeferredPromise(); + const request = createDeferredPromise(); + + const activeRequest = manager.withProviderMaintenanceRuntime( + { dataDir }, + async () => { + requestStarted.resolve(); + return request.promise; + }, + ); + await requestStarted.promise; + await manager.replaceBaseShellEnv({ PATH: "/new/bin:/usr/bin" }); + await manager.shutdownAll(); + + expect(runtime.shutdown).toHaveBeenCalledOnce(); + request.resolve(); + await activeRequest; + expect(runtime.shutdown).toHaveBeenCalledOnce(); + }); + + it("keeps runtimes when only fnm multishell PATH entries change", async () => { + const dataDir = await makeTempDir("bb-provider-maintenance-fnm-"); + const runtime = createFakeRuntime(); + const createRuntime = vi.fn(() => runtime); + const manager = new RuntimeManager({ + createRuntime, + shellEnv: { + PATH: "/Users/me/Library/Caches/fnm_multishells/123_1/bin:/usr/bin", + }, + }); + + await manager.ensureProviderMaintenanceRuntime({ dataDir }); + await manager.replaceBaseShellEnv({ + PATH: "/Users/me/Library/Caches/fnm_multishells/456_2/bin:/usr/bin", + }); + + expect(runtime.shutdown).not.toHaveBeenCalled(); + await expect( + manager.ensureProviderMaintenanceRuntime({ dataDir }), + ).resolves.toBe(runtime); + expect(createRuntime).toHaveBeenCalledOnce(); + + await manager.replaceBaseShellEnv({ + PATH: "/Users/me/Library/Caches/fnm_multishells/789_3/bin:/opt/bin", + }); + expect(runtime.shutdown).toHaveBeenCalledOnce(); + }); + it("does not let stale provider maintenance creation replace a newer runtime", async () => { const dataDir = await makeTempDir("bb-provider-maintenance-race-"); const staleRuntime = createFakeRuntime(); @@ -1188,13 +1283,14 @@ describe("RuntimeManager", () => { }); await expect(currentRuntimePromise).resolves.toBe(currentRuntime); - expect(managerInternals.providerMaintenanceRuntime).toBe(currentRuntime); staleCreation.resolve(staleRuntime); await expect(staleRuntimePromise).resolves.toBe(staleRuntime); await replaceShellEnvPromise; - expect(managerInternals.providerMaintenanceRuntime).toBe(currentRuntime); + await expect( + manager.ensureProviderMaintenanceRuntime({ dataDir }), + ).resolves.toBe(currentRuntime); expect(staleRuntime.shutdown).toHaveBeenCalledTimes(1); expect(currentRuntime.shutdown).not.toHaveBeenCalled(); }); @@ -1492,7 +1588,7 @@ describe("RuntimeManager", () => { provisionWorkspace, createRuntime: () => runtime, shellEnv: { - PATH: "/tmp/fnm_multishells/first/bin:/usr/bin", + PATH: "/tmp/first/bin:/usr/bin", }, }); @@ -1503,7 +1599,7 @@ describe("RuntimeManager", () => { runtime.setPendingTurnStart("thread-1", true); await manager.replaceBaseShellEnv({ - PATH: "/tmp/fnm_multishells/second/bin:/usr/bin", + PATH: "/tmp/second/bin:/usr/bin", }); expect(manager.get("env-1")?.runtime).toBe(runtime); @@ -1511,7 +1607,7 @@ describe("RuntimeManager", () => { runtime.setPendingTurnStart("thread-1", false); await manager.replaceBaseShellEnv({ - PATH: "/tmp/fnm_multishells/third/bin:/usr/bin", + PATH: "/tmp/third/bin:/usr/bin", }); expect(manager.get("env-1")).toBeUndefined(); diff --git a/apps/host-daemon/src/runtime-manager.ts b/apps/host-daemon/src/runtime-manager.ts index f3409b51783..e34cf139449 100644 --- a/apps/host-daemon/src/runtime-manager.ts +++ b/apps/host-daemon/src/runtime-manager.ts @@ -226,9 +226,19 @@ interface PendingEnvironmentProvision { done: Promise; } -interface PendingProviderMaintenanceRuntime { - generation: number; - promise: Promise; +interface ProviderMaintenanceRuntimeSlot { + runtime: Promise; + activeRequests: number; +} + +const FNM_MULTISHELL_PATH_ENTRY = /([\\/]fnm_multishells[\\/])[^\\/]+/; + +function comparableShellEnvValue(key: string, value: string): string { + if (key !== "PATH") return value; + return value + .split(path.delimiter) + .map((entry) => entry.replace(FNM_MULTISHELL_PATH_ENTRY, "$1*")) + .join(path.delimiter); } function shellEnvEquals( @@ -240,7 +250,14 @@ function shellEnvEquals( if (leftEntries.length !== rightEntries.length) { return false; } - return leftEntries.every(([key, value]) => right[key] === value); + return leftEntries.every(([key, value]) => { + const rightValue = right[key]; + return ( + rightValue !== undefined && + comparableShellEnvValue(key, value) === + comparableShellEnvValue(key, rightValue) + ); + }); } function providerProcessEnvFromShellEnv( @@ -287,11 +304,9 @@ export class RuntimeManager { Map>> >(); private readonly threadControlTails = new Map>(); - private providerMaintenanceRuntime: AgentRuntime | null = null; - private pendingProviderMaintenanceRuntime: PendingProviderMaintenanceRuntime | null = - null; - private providerMaintenanceRuntimeGeneration = 0; - private providerMaintenanceActiveRequests = 0; + private providerMaintenanceSlot: ProviderMaintenanceRuntimeSlot | null = null; + private readonly retiredProviderMaintenanceSlots = + new Set(); private providerMaintenanceIdleTimer: ReturnType | null = null; readonly providerInstallationGate: ProviderInstallationGate; @@ -570,7 +585,7 @@ export class RuntimeManager { this.baseShellEnv = { ...shellEnv }; this.shellEnvGeneration += 1; this.providerInstallationGate.clear(); - await this.shutdownProviderMaintenanceRuntime(); + await this.retireProviderMaintenanceRuntime(); await this.evictIdleRuntimeEntries(); } @@ -751,7 +766,7 @@ export class RuntimeManager { async invalidateProviderMaintenanceRuntime(): Promise { this.providerInstallationGate.clear(); try { - await this.shutdownProviderMaintenanceRuntime(); + await this.retireProviderMaintenanceRuntime(); } catch (error) { this.options.logger?.warn( { err: error }, @@ -760,29 +775,35 @@ export class RuntimeManager { } } - private async shutdownProviderMaintenanceRuntime(): Promise { + private async retireProviderMaintenanceRuntime(): Promise { this.clearProviderMaintenanceIdleTimer(); - const existingRuntime = this.providerMaintenanceRuntime; - const pendingRuntime = this.pendingProviderMaintenanceRuntime; - this.providerMaintenanceRuntimeGeneration += 1; - this.providerMaintenanceRuntime = null; - if (this.pendingProviderMaintenanceRuntime === pendingRuntime) { - this.pendingProviderMaintenanceRuntime = null; + const slot = this.providerMaintenanceSlot; + this.providerMaintenanceSlot = null; + if (!slot) return; + if (slot.activeRequests > 0) { + this.retiredProviderMaintenanceSlots.add(slot); + return; } + await this.shutdownProviderMaintenanceSlot(slot); + } - const resolvedPendingRuntime = pendingRuntime - ? await pendingRuntime.promise.catch(() => null) - : null; - if ( - resolvedPendingRuntime && - this.providerMaintenanceRuntime === resolvedPendingRuntime - ) { - this.providerMaintenanceRuntime = null; - } + private async shutdownProviderMaintenanceSlot( + slot: ProviderMaintenanceRuntimeSlot, + ): Promise { + const runtime = await slot.runtime.catch(() => null); + await runtime?.shutdown(); + } - const runtimes = [...new Set([existingRuntime, resolvedPendingRuntime])]; + private async shutdownAllProviderMaintenanceRuntimes(): Promise { + this.clearProviderMaintenanceIdleTimer(); + const slots = [ + this.providerMaintenanceSlot, + ...this.retiredProviderMaintenanceSlots, + ].filter((slot) => slot !== null); + this.providerMaintenanceSlot = null; + this.retiredProviderMaintenanceSlots.clear(); await Promise.all( - runtimes.map((runtime) => runtime?.shutdown() ?? Promise.resolve()), + slots.map((slot) => this.shutdownProviderMaintenanceSlot(slot)), ); } @@ -792,23 +813,30 @@ export class RuntimeManager { this.providerMaintenanceIdleTimer = null; } - private scheduleProviderMaintenanceIdleShutdown(): void { - this.clearProviderMaintenanceIdleTimer(); - if ( - this.providerMaintenanceActiveRequests > 0 || - (this.providerMaintenanceRuntime === null && - this.pendingProviderMaintenanceRuntime === null) - ) { + private releaseIdleProviderMaintenanceSlot( + slot: ProviderMaintenanceRuntimeSlot, + ): void { + if (this.retiredProviderMaintenanceSlots.delete(slot)) { + void this.shutdownProviderMaintenanceSlot(slot).catch((error) => { + this.options.logger?.warn( + { err: error }, + "Failed to shut down retired provider maintenance runtime", + ); + }); return; } + if (this.providerMaintenanceSlot !== slot) return; + this.clearProviderMaintenanceIdleTimer(); const timeoutMs = this.options.providerMaintenanceIdleTimeoutMs ?? PROVIDER_MAINTENANCE_IDLE_TIMEOUT_MS; this.providerMaintenanceIdleTimer = setTimeout(() => { this.providerMaintenanceIdleTimer = null; - if (this.providerMaintenanceActiveRequests > 0) return; - void this.shutdownProviderMaintenanceRuntime().catch((error) => { + if (this.providerMaintenanceSlot !== slot || slot.activeRequests > 0) { + return; + } + void this.retireProviderMaintenanceRuntime().catch((error) => { this.options.logger?.warn( { err: error }, "Failed to shut down idle provider maintenance runtime", @@ -831,40 +859,31 @@ export class RuntimeManager { await this.cleanupUnusedInjectedSkillStagingDirs([]); } - async ensureProviderMaintenanceRuntime(args: { + private currentProviderMaintenanceSlot(args: { dataDir: string; - }): Promise { - if (this.providerMaintenanceRuntime) { - return this.providerMaintenanceRuntime; + }): ProviderMaintenanceRuntimeSlot { + if (this.providerMaintenanceSlot) { + return this.providerMaintenanceSlot; } - if (this.pendingProviderMaintenanceRuntime) { - return this.pendingProviderMaintenanceRuntime.promise; - } - - const generation = this.providerMaintenanceRuntimeGeneration; - let pendingRuntime!: PendingProviderMaintenanceRuntime; - const promise = Promise.resolve() - .then(() => this.createProviderMaintenanceRuntime(args)) - .then((runtime) => { - if ( - this.pendingProviderMaintenanceRuntime === pendingRuntime && - this.providerMaintenanceRuntimeGeneration === generation - ) { - this.providerMaintenanceRuntime = runtime; - } - return runtime; - }) - .finally(() => { - if (this.pendingProviderMaintenanceRuntime === pendingRuntime) { - this.pendingProviderMaintenanceRuntime = null; - } - }); - pendingRuntime = { - generation, - promise, + const slot: ProviderMaintenanceRuntimeSlot = { + runtime: Promise.resolve().then(() => + this.createProviderMaintenanceRuntime(args, slot), + ), + activeRequests: 0, }; - this.pendingProviderMaintenanceRuntime = pendingRuntime; - return promise; + slot.runtime.catch(() => { + if (this.providerMaintenanceSlot === slot) { + this.providerMaintenanceSlot = null; + } + }); + this.providerMaintenanceSlot = slot; + return slot; + } + + async ensureProviderMaintenanceRuntime(args: { + dataDir: string; + }): Promise { + return this.currentProviderMaintenanceSlot(args).runtime; } async withProviderMaintenanceRuntime( @@ -872,14 +891,14 @@ export class RuntimeManager { request: (runtime: AgentRuntime) => Promise, ): Promise { this.clearProviderMaintenanceIdleTimer(); - this.providerMaintenanceActiveRequests += 1; + const slot = this.currentProviderMaintenanceSlot(args); + slot.activeRequests += 1; try { - const runtime = await this.ensureProviderMaintenanceRuntime(args); - return await request(runtime); + return await request(await slot.runtime); } finally { - this.providerMaintenanceActiveRequests -= 1; - if (this.providerMaintenanceActiveRequests === 0) { - this.scheduleProviderMaintenanceIdleShutdown(); + slot.activeRequests -= 1; + if (slot.activeRequests === 0) { + this.releaseIdleProviderMaintenanceSlot(slot); } } } @@ -1096,7 +1115,7 @@ export class RuntimeManager { for (const entry of entries) { await entry.runtime.shutdown(); } - await this.shutdownProviderMaintenanceRuntime(); + await this.shutdownAllProviderMaintenanceRuntimes(); await this.stopWatchingDataDirSkillsRoot(); this.stopWatchingDataDirSkillsRoot = STOP_WATCHING; await this.cleanupUnusedInjectedSkillStagingDirs([]); @@ -1149,9 +1168,10 @@ export class RuntimeManager { return events; } - private async createProviderMaintenanceRuntime(args: { - dataDir: string; - }): Promise { + private async createProviderMaintenanceRuntime( + args: { dataDir: string }, + slot: ProviderMaintenanceRuntimeSlot, + ): Promise { const workspacePath = path.join( args.dataDir, PROVIDER_MAINTENANCE_WORKSPACE_DIR, @@ -1185,10 +1205,10 @@ export class RuntimeManager { onProcessExit: (info) => { if ( runtime && - this.providerMaintenanceRuntime === runtime && + this.providerMaintenanceSlot === slot && runtime.listRunningProviders().length === 0 ) { - this.providerMaintenanceRuntime = null; + this.providerMaintenanceSlot = null; } this.options.onProcessExit?.(info); },