Keep provider usage requests alive across login-shell env refreshes

apps/host-daemon/src/runtime-manager.test.ts

@@ -54,10 +54,7 @@ interface RuntimeOptionsRef {
 }
 
 interface RuntimeManagerProviderMaintenanceInternals {
-  createProviderMaintenanceRuntime: (args: {
-    dataDir: string;
-  }) => Promise<AgentRuntime>;
-  providerMaintenanceRuntime: AgentRuntime | null;
+  createProviderMaintenanceRuntime: () => Promise<AgentRuntime>;
 }
 
 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<string>();
+    const usageStarted = createDeferredPromise<AgentRuntime>();
+
+    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<void>();
+    const request = createDeferredPromise<void>();
+
+    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,15 +1599,15 @@ 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);
     expect(runtime.shutdown).not.toHaveBeenCalled();
 
     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();