From 2b539e46d2858fa069629f51bf362181cea6163f Mon Sep 17 00:00:00 2001 From: Michael Yong Date: Mon, 3 Aug 2026 11:35:52 -0700 Subject: [PATCH] Invalidate Codex model listing after CLI update --- .../provider-cli-install-store.ts | 2 + .../provider-cli-install.test.tsx | 9 +- .../cache-owners/system-cache-effects.ts | 16 ++ apps/host-daemon/src/app.ts | 8 +- .../src/command-dispatch-support.ts | 4 + apps/host-daemon/src/command-dispatch.test.ts | 163 +++++++++++++++++- apps/host-daemon/src/command-dispatch.ts | 62 +++++-- apps/host-daemon/src/runtime-manager.test.ts | 39 +++++ apps/host-daemon/src/runtime-manager.ts | 90 +++++++++- 9 files changed, 366 insertions(+), 27 deletions(-) diff --git a/apps/app/src/components/provider-cli/provider-cli-install-store.ts b/apps/app/src/components/provider-cli/provider-cli-install-store.ts index 835760210..2f28ccaee 100644 --- a/apps/app/src/components/provider-cli/provider-cli-install-store.ts +++ b/apps/app/src/components/provider-cli/provider-cli-install-store.ts @@ -8,6 +8,7 @@ import type { ProviderCliInstallLogDialogState } from "@/components/dialogs/Prov import type { ProviderCliActionableIssue } from "@/components/provider-cli/provider-cli-install"; import { appToast } from "@/components/ui/app-toast"; import { invalidateHostProviderCliStatus } from "@/hooks/cache-owners/provider-cli-status-cache-owner"; +import { invalidateSystemExecutionOptions } from "@/hooks/cache-owners/system-cache-effects"; import { sdk } from "@/lib/sdk"; type ProviderCliInstallCompletedEvent = Extract< @@ -202,6 +203,7 @@ function runInstall(job: ProviderCliInstallJob): void { if (completedEvent?.success) { if (queryClient !== null) { void invalidateHostProviderCliStatus({ queryClient, hostId }); + void invalidateSystemExecutionOptions({ queryClient, hostId }); } return; } diff --git a/apps/app/src/components/provider-cli/provider-cli-install.test.tsx b/apps/app/src/components/provider-cli/provider-cli-install.test.tsx index 7b31f88a3..8d5c951c5 100644 --- a/apps/app/src/components/provider-cli/provider-cli-install.test.tsx +++ b/apps/app/src/components/provider-cli/provider-cli-install.test.tsx @@ -10,7 +10,10 @@ import type { import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { sdk } from "@/lib/sdk"; import { appToast } from "@/components/ui/app-toast"; -import { hostProviderCliStatusQueryKey } from "@/hooks/queries/query-keys"; +import { + allSystemExecutionOptionsQueryKeyPrefix, + hostProviderCliStatusQueryKey, +} from "@/hooks/queries/query-keys"; import type { ProviderCliActionableIssue } from "./provider-cli-install"; import { useProviderCliInstallRunner } from "./provider-cli-install"; import { @@ -203,6 +206,10 @@ describe("useProviderCliInstallRunner", () => { expect(invalidateQueries).toHaveBeenCalledWith({ queryKey: hostProviderCliStatusQueryKey("host_1"), }); + expect(invalidateQueries).toHaveBeenCalledWith({ + queryKey: allSystemExecutionOptionsQueryKeyPrefix(), + predicate: expect.any(Function), + }); expect(appToastMock.success).not.toHaveBeenCalled(); }); diff --git a/apps/app/src/hooks/cache-owners/system-cache-effects.ts b/apps/app/src/hooks/cache-owners/system-cache-effects.ts index a24d477f9..453e80af1 100644 --- a/apps/app/src/hooks/cache-owners/system-cache-effects.ts +++ b/apps/app/src/hooks/cache-owners/system-cache-effects.ts @@ -49,6 +49,10 @@ interface SystemExecutionOptionsCacheArgs extends QueryClientArg { providerId: string | null; } +interface SystemExecutionOptionsInvalidationArgs extends QueryClientArg { + hostId: string; +} + export function seedSystemExecutionOptionsCache({ environmentId, executionOptions, @@ -134,6 +138,18 @@ export function invalidateSystemConfig({ queryClient }: QueryClientArg): void { queryClient.invalidateQueries({ queryKey: systemConfigQueryKey() }); } +/** Refresh provider/model catalogs after a provider CLI install or update. */ +export function invalidateSystemExecutionOptions({ + hostId, + queryClient, +}: SystemExecutionOptionsInvalidationArgs): Promise { + return queryClient.invalidateQueries({ + queryKey: allSystemExecutionOptionsQueryKeyPrefix(), + predicate: (query) => + query.queryKey[2] === hostId || query.queryKey[2] === null, + }); +} + /** Refresh settings and timeline projections after a General settings write. */ export function invalidateGeneralSettingsDependencies({ queryClient, diff --git a/apps/host-daemon/src/app.ts b/apps/host-daemon/src/app.ts index 7ce7dd1bc..68e6dbb50 100644 --- a/apps/host-daemon/src/app.ts +++ b/apps/host-daemon/src/app.ts @@ -747,10 +747,10 @@ export async function createHostDaemonApp( terminalManager, listModels: async (args) => { await refreshRuntimeShellEnv(); - const runtime = await runtimeManager.ensureProviderMaintenanceRuntime({ - dataDir: options.dataDir, - }); - return runtime.listModels(args); + return runtimeManager.withProviderMaintenanceRuntime( + { dataDir: options.dataDir }, + (runtime) => runtime.listModels(args), + ); }, resolveInteractiveRequest: async (request) => { interactiveRequestRegistry.resolve(request); diff --git a/apps/host-daemon/src/command-dispatch-support.ts b/apps/host-daemon/src/command-dispatch-support.ts index 3fd4708ff..06979cb8d 100644 --- a/apps/host-daemon/src/command-dispatch-support.ts +++ b/apps/host-daemon/src/command-dispatch-support.ts @@ -12,6 +12,7 @@ import type { HostDaemonInjectedSkillSource, HostDaemonOnlineRpcCommand, HostDaemonConnectTunnelIdentity, + ProviderCliInstallRequest, ProviderCliStatus, WorkspaceContext, } from "@bb/host-daemon-contract"; @@ -57,6 +58,9 @@ export interface CommandDispatchOptions { getProviderCliStatusForProvider?: ( providerId: string, ) => Promise; + streamProviderCliInstall?: ( + args: ProviderCliInstallRequest & { env?: NodeJS.ProcessEnv }, + ) => ReadableStream; resolveInteractiveRequest?: ( request: InteractiveResolveCommandInput, ) => Promise; diff --git a/apps/host-daemon/src/command-dispatch.test.ts b/apps/host-daemon/src/command-dispatch.test.ts index 790f07c9c..a1845b2ff 100644 --- a/apps/host-daemon/src/command-dispatch.test.ts +++ b/apps/host-daemon/src/command-dispatch.test.ts @@ -4,11 +4,15 @@ import path from "node:path"; import type { AgentRuntime } from "@bb/agent-runtime"; import type { HostDaemonInjectedSkillSource, + ProviderCliInstallEvent, ProviderCliStatus, } from "@bb/host-daemon-contract"; import type { HostWorkspace } from "@bb/host-workspace"; import { afterEach, describe, expect, it, vi, type Mock } from "vitest"; -import { dispatchCommand } from "./command-dispatch.js"; +import { + dispatchCommand, + dispatchOnlineRpcCommand, +} from "./command-dispatch.js"; import type { CommandOf } from "./command-dispatch-support.js"; import { RuntimeManager } from "./runtime-manager.js"; @@ -215,6 +219,20 @@ function createRuntime(): FakeDispatchRuntime { }; } +function createProviderCliInstallEventStream( + events: readonly ProviderCliInstallEvent[], +): ReadableStream { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + for (const event of events) { + controller.enqueue(encoder.encode(`${JSON.stringify(event)}\n`)); + } + controller.close(); + }, + }); +} + describe("dispatchCommand", () => { it("flushes buffered events before reporting thread.stop success", async () => { const runtime = createRuntime(); @@ -595,6 +613,149 @@ describe("dispatchCommand", () => { expect(runtime.startThread).toHaveBeenCalledOnce(); }); + it("invalidates the provider maintenance runtime after a successful Codex CLI update", async () => { + const dataDir = await makeTempDir("bb-command-dispatch-provider-cli-"); + const staleRuntime = createRuntime(); + const freshRuntime = createRuntime(); + const createRuntimeSpy = vi.fn(() => staleRuntime); + createRuntimeSpy.mockReturnValueOnce(staleRuntime); + createRuntimeSpy.mockReturnValueOnce(freshRuntime); + const manager = new RuntimeManager({ + createRuntime: createRuntimeSpy, + dataDir, + provisionWorkspace: async () => createWorkspace(), + }); + await manager.ensureProviderMaintenanceRuntime({ dataDir }); + + const events: ProviderCliInstallEvent[] = [ + { + type: "started", + provider: "codex", + command: "codex update", + }, + { + type: "completed", + provider: "codex", + exitCode: 0, + signal: null, + success: true, + }, + ]; + const streamProviderCliInstall = vi.fn(() => + createProviderCliInstallEventStream(events), + ); + const command: CommandOf<"provider_cli.install"> = { + type: "provider_cli.install", + provider: "codex", + actionKind: "update", + }; + + const result = await dispatchOnlineRpcCommand(command, { + dataDir, + eventSink: { + emit: vi.fn(), + flush: vi.fn(async () => undefined), + }, + fetchProjectAttachment: async () => { + throw new Error("Unexpected project attachment fetch"); + }, + runtimeManager: manager, + streamProviderCliInstall, + threadStorageRootPath: "/tmp/bb-thread-storage", + }); + + expect(result).toEqual({ events }); + expect(streamProviderCliInstall).toHaveBeenCalledWith( + expect.objectContaining({ + actionKind: "update", + provider: "codex", + }), + ); + expect(staleRuntime.shutdown).toHaveBeenCalledOnce(); + + await manager.ensureProviderMaintenanceRuntime({ dataDir }); + expect(createRuntimeSpy).toHaveBeenCalledTimes(2); + expect(freshRuntime.shutdown).not.toHaveBeenCalled(); + }); + + it("keeps the provider maintenance runtime after failed or non-Codex CLI installs", async () => { + const cases: Array<{ + actionKind: CommandOf<"provider_cli.install">["actionKind"]; + events: ProviderCliInstallEvent[]; + provider: CommandOf<"provider_cli.install">["provider"]; + }> = [ + { + actionKind: "update", + provider: "codex", + events: [ + { + type: "completed", + provider: "codex", + exitCode: 1, + signal: null, + success: false, + }, + ], + }, + { + actionKind: "update", + provider: "claudeCode", + events: [ + { + type: "completed", + provider: "claudeCode", + exitCode: 0, + signal: null, + success: true, + }, + ], + }, + ]; + + for (const testCase of cases) { + const dataDir = await makeTempDir("bb-command-dispatch-provider-cli-"); + const runtime = createRuntime(); + const createRuntimeSpy = vi.fn(() => runtime); + const manager = new RuntimeManager({ + createRuntime: createRuntimeSpy, + dataDir, + provisionWorkspace: async () => createWorkspace(), + }); + await manager.ensureProviderMaintenanceRuntime({ dataDir }); + const streamProviderCliInstall = vi.fn(() => + createProviderCliInstallEventStream(testCase.events), + ); + + const result = await dispatchOnlineRpcCommand( + { + type: "provider_cli.install", + provider: testCase.provider, + actionKind: testCase.actionKind, + }, + { + dataDir, + eventSink: { + emit: vi.fn(), + flush: vi.fn(async () => undefined), + }, + fetchProjectAttachment: async () => { + throw new Error("Unexpected project attachment fetch"); + }, + runtimeManager: manager, + streamProviderCliInstall, + threadStorageRootPath: "/tmp/bb-thread-storage", + }, + ); + + expect(result).toEqual({ events: testCase.events }); + expect(runtime.shutdown).not.toHaveBeenCalled(); + await expect( + manager.ensureProviderMaintenanceRuntime({ dataDir }), + ).resolves.toBe(runtime); + expect(createRuntimeSpy).toHaveBeenCalledTimes(1); + } + }); + // Regression: a thread.start whose freshly staged skill catalog differed // from the busy runtime's catalog used to fail the command (and brick the // thread) instead of reusing the runtime. This drives the real plumbing — diff --git a/apps/host-daemon/src/command-dispatch.ts b/apps/host-daemon/src/command-dispatch.ts index 310ecdaff..e46f00570 100644 --- a/apps/host-daemon/src/command-dispatch.ts +++ b/apps/host-daemon/src/command-dispatch.ts @@ -176,15 +176,24 @@ async function installProviderCliOnHost( const env = providerCliEnvFromShellEnv( options.runtimeManager.getShellEnv(), ); - return { - events: await readProviderCliInstallEvents( - streamProviderCliInstall({ - provider: command.provider, - actionKind: command.actionKind, - env, - }), - ), - }; + const streamInstall = + options.streamProviderCliInstall ?? streamProviderCliInstall; + const events = await readProviderCliInstallEvents( + streamInstall({ + provider: command.provider, + actionKind: command.actionKind, + env, + }), + ); + if ( + shouldInvalidateProviderMaintenanceRuntimeAfterProviderCliInstall({ + command, + events, + }) + ) { + await options.runtimeManager.invalidateProviderMaintenanceRuntime(); + } + return { events }; } catch (error) { if (error instanceof ProviderCliInstallInProgressError) { return { @@ -201,6 +210,23 @@ async function installProviderCliOnHost( } } +function shouldInvalidateProviderMaintenanceRuntimeAfterProviderCliInstall(args: { + command: CommandOf<"provider_cli.install">; + events: readonly ProviderCliInstallEvent[]; +}): boolean { + return ( + // Codex model listing goes through the resident provider-maintenance + // app-server, so a Codex CLI update can leave a stale model catalog alive. + args.command.provider === "codex" && + args.events.some( + (event) => + event.type === "completed" && + event.provider === args.command.provider && + event.success, + ) + ); +} + const commandHandlers: CommandHandlerMap = { "thread.start": async (command, options) => { const release = options.runtimeManager.retainEnvironmentForThreadCommand( @@ -302,15 +328,17 @@ const commandHandlers: CommandHandlerMap = { return {}; }, "thread.unarchive": async (command, options) => { - const runtime = - await options.runtimeManager.ensureProviderMaintenanceRuntime({ + await options.runtimeManager.withProviderMaintenanceRuntime( + { dataDir: options.dataDir, - }); - await runtime.unarchiveThread({ - threadId: command.threadId, - providerId: command.providerId, - providerThreadId: command.providerThreadId, - }); + }, + (runtime) => + runtime.unarchiveThread({ + threadId: command.threadId, + providerId: command.providerId, + providerThreadId: command.providerThreadId, + }), + ); return {}; }, "interactive.resolve": resolveInteractiveRequest, diff --git a/apps/host-daemon/src/runtime-manager.test.ts b/apps/host-daemon/src/runtime-manager.test.ts index d26eace17..4c9f771e7 100644 --- a/apps/host-daemon/src/runtime-manager.test.ts +++ b/apps/host-daemon/src/runtime-manager.test.ts @@ -1222,6 +1222,45 @@ describe("RuntimeManager", () => { ); }); + it("defers stale provider maintenance runtime shutdown until active maintenance work finishes", async () => { + const dataDir = await makeTempDir("bb-provider-maintenance-lease-"); + const firstRuntime = createFakeRuntime(); + const secondRuntime = createFakeRuntime(); + const activeWork = createDeferred(); + const createRuntime = vi + .fn() + .mockReturnValueOnce(firstRuntime) + .mockReturnValueOnce(secondRuntime); + const manager = new RuntimeManager({ createRuntime }); + let leasedRuntime: AgentRuntime | null = null; + + const workPromise = manager.withProviderMaintenanceRuntime( + { dataDir }, + async (runtime) => { + leasedRuntime = runtime; + await activeWork.promise; + return "done"; + }, + ); + + await vi.waitFor(() => { + expect(leasedRuntime).toBe(firstRuntime); + }); + + await manager.invalidateProviderMaintenanceRuntime(); + expect(firstRuntime.shutdown).not.toHaveBeenCalled(); + await expect( + manager.ensureProviderMaintenanceRuntime({ dataDir }), + ).resolves.toBe(secondRuntime); + + activeWork.resolve(undefined); + await expect(workPromise).resolves.toBe("done"); + await vi.waitFor(() => { + expect(firstRuntime.shutdown).toHaveBeenCalledTimes(1); + }); + expect(secondRuntime.shutdown).not.toHaveBeenCalled(); + }); + it("does not let stale provider maintenance creation replace a newer runtime", async () => { const dataDir = await makeTempDir("bb-provider-maintenance-race-"); const staleRuntime = createFakeRuntime(); diff --git a/apps/host-daemon/src/runtime-manager.ts b/apps/host-daemon/src/runtime-manager.ts index 8bd751b99..6a39ba3ef 100644 --- a/apps/host-daemon/src/runtime-manager.ts +++ b/apps/host-daemon/src/runtime-manager.ts @@ -279,6 +279,8 @@ export class RuntimeManager { private pendingProviderMaintenanceRuntime: PendingProviderMaintenanceRuntime | null = null; private providerMaintenanceRuntimeGeneration = 0; + private providerMaintenanceRuntimeLeaseCount = 0; + private readonly retiredProviderMaintenanceRuntimes = new Set(); private managedShellEnv: NonNullable = {}; private stopWatchingDataDirSkillsRoot: StopWatching = STOP_WATCHING; @@ -610,6 +612,86 @@ export class RuntimeManager { this.managedShellEnv = { ...shellEnv }; } + async invalidateProviderMaintenanceRuntime(): Promise { + try { + await this.shutdownProviderMaintenanceRuntime(); + } catch (error) { + this.options.logger?.warn( + { err: error }, + "Failed to shut down provider maintenance runtime during invalidation", + ); + } + } + + async withProviderMaintenanceRuntime( + args: { dataDir: string }, + work: (runtime: AgentRuntime) => Promise, + ): Promise { + const release = this.retainProviderMaintenanceRuntime(); + try { + const runtime = await this.ensureProviderMaintenanceRuntime(args); + return await work(runtime); + } finally { + release(); + } + } + + private retainProviderMaintenanceRuntime(): () => void { + this.providerMaintenanceRuntimeLeaseCount += 1; + let released = false; + return () => { + if (released) { + return; + } + released = true; + this.providerMaintenanceRuntimeLeaseCount -= 1; + if (this.providerMaintenanceRuntimeLeaseCount === 0) { + void this.shutdownRetiredProviderMaintenanceRuntimes().catch( + (error: unknown) => { + this.options.logger?.warn( + { err: error }, + "Failed to shut down retired provider maintenance runtime", + ); + }, + ); + } + }; + } + + private async shutdownRetiredProviderMaintenanceRuntimes(): Promise { + if ( + this.providerMaintenanceRuntimeLeaseCount > 0 || + this.retiredProviderMaintenanceRuntimes.size === 0 + ) { + return; + } + + const runtimes = [...this.retiredProviderMaintenanceRuntimes]; + this.retiredProviderMaintenanceRuntimes.clear(); + await Promise.all(runtimes.map((runtime) => runtime.shutdown())); + } + + private async retireProviderMaintenanceRuntimes( + runtimes: readonly (AgentRuntime | null)[], + ): Promise { + const uniqueRuntimes = [ + ...new Set( + runtimes.filter((runtime): runtime is AgentRuntime => runtime !== null), + ), + ]; + if (uniqueRuntimes.length === 0) { + return; + } + if (this.providerMaintenanceRuntimeLeaseCount > 0) { + for (const runtime of uniqueRuntimes) { + this.retiredProviderMaintenanceRuntimes.add(runtime); + } + return; + } + + await Promise.all(uniqueRuntimes.map((runtime) => runtime.shutdown())); + } + private async shutdownProviderMaintenanceRuntime(): Promise { const existingRuntime = this.providerMaintenanceRuntime; const pendingRuntime = this.pendingProviderMaintenanceRuntime; @@ -629,10 +711,10 @@ export class RuntimeManager { this.providerMaintenanceRuntime = null; } - const runtimes = [...new Set([existingRuntime, resolvedPendingRuntime])]; - await Promise.all( - runtimes.map((runtime) => runtime?.shutdown() ?? Promise.resolve()), - ); + await this.retireProviderMaintenanceRuntimes([ + existingRuntime, + resolvedPendingRuntime, + ]); } private async evictIdleRuntimeEntries(): Promise {