diff --git a/src/agentMode/index.ts b/src/agentMode/index.ts index 3f9ffb69..508bd989 100644 --- a/src/agentMode/index.ts +++ b/src/agentMode/index.ts @@ -132,18 +132,19 @@ export function createAgentSessionManager(app: App, plugin: CopilotPlugin): Agen const settings = getSettings(); if (!isAgentModeEnabled(settings)) return manager; - const preloads: Promise[] = []; + // Per-backend preload registration: each backend's status flips + // independently. The chat UI gates on the active backend's status; the + // picker reads every backend's status to render per-backend loading rows. for (const descriptor of listBackendDescriptors()) { if (descriptor.getInstallState(settings).kind !== "ready") continue; - preloads.push( - manager - .preloadModels(descriptor.id) - .catch((e) => logError(`[AgentMode] preload ${descriptor.id} failed`, e)) + const promise = manager.preloadModels(descriptor.id); + manager.registerPreload( + descriptor.id, + promise.catch((e) => { + logError(`[AgentMode] preload ${descriptor.id} failed`, e); + throw e; + }) ); } - // Aggregate so the chat UI can gate its first render until every - // backend's catalog has settled — the model picker should never flash - // empty before the cache populates. - manager.setPreloadPromise(Promise.allSettled(preloads).then(() => undefined)); return manager; } diff --git a/src/agentMode/session/AgentModelPreloader.test.ts b/src/agentMode/session/AgentModelPreloader.test.ts new file mode 100644 index 00000000..96a91697 --- /dev/null +++ b/src/agentMode/session/AgentModelPreloader.test.ts @@ -0,0 +1,213 @@ +/** + * Tests the warm-reuse contract added so the first chat-open per backend + * doesn't pay a second subprocess spawn + initialize handshake. The probe + * subprocess survives `runProbe` and is handed to the manager via + * `takeWarm`; failures and shutdowns must dispose the still-warm proc. + */ +import { App, FileSystemAdapter } from "obsidian"; +import type CopilotPlugin from "@/main"; +import { AgentModelPreloader } from "./AgentModelPreloader"; +import type { BackendDescriptor, BackendProcess, BackendState, ModelSelection } from "./types"; + +jest.mock("@/logger", () => ({ + logInfo: jest.fn(), + logWarn: jest.fn(), + logError: jest.fn(), +})); + +jest.mock("@/settings/model", () => ({ + getSettings: jest.fn(() => ({ agentMode: {} })), +})); + +function buildApp(): App { + const adapter = new (FileSystemAdapter as unknown as new (basePath: string) => unknown)("/vault"); + return { vault: { adapter } } as unknown as App; +} + +function buildPlugin(): CopilotPlugin { + return { manifest: { version: "1.0.0" } } as unknown as CopilotPlugin; +} + +interface MockProcHandle { + proc: BackendProcess; + start: jest.Mock; + shutdown: jest.Mock; + newSession: jest.Mock; + exitListeners: Set<() => void>; + emitExit: () => void; +} + +function makeMockProc(opts?: { newSessionState?: BackendState }): MockProcHandle { + const start = jest.fn(async () => undefined); + const shutdown = jest.fn(async () => undefined); + const state: BackendState = opts?.newSessionState ?? { + model: { + current: { baseModelId: "claude-sonnet", effort: null } satisfies ModelSelection, + availableModels: [ + { + baseModelId: "claude-sonnet", + name: "Claude Sonnet", + provider: "anthropic", + effortOptions: [], + }, + ], + }, + mode: null, + }; + const newSession = jest.fn(async () => ({ sessionId: "probe-1", state })); + const exitListeners = new Set<() => void>(); + const proc: BackendProcess = { + start, + isRunning: () => true, + onExit: (fn) => { + exitListeners.add(fn); + return () => exitListeners.delete(fn); + }, + setPermissionPrompter: jest.fn(), + registerSessionHandler: jest.fn(() => () => {}), + newSession, + prompt: jest.fn(), + cancel: jest.fn(), + setSessionModel: jest.fn(), + isSetSessionModelSupported: () => true, + setSessionMode: jest.fn(), + isSetSessionModeSupported: () => true, + setSessionConfigOption: jest.fn(), + isSetSessionConfigOptionSupported: () => true, + listSessions: jest.fn(), + resumeSession: jest.fn(), + loadSession: jest.fn(), + supportsMcpTransport: () => false, + shutdown, + }; + return { + proc, + start, + shutdown, + newSession, + exitListeners, + emitExit: () => { + for (const fn of exitListeners) fn(); + }, + }; +} + +function buildDescriptor(makeProc: () => MockProcHandle): { + descriptor: BackendDescriptor; + procHandle: MockProcHandle; +} { + const procHandle = makeProc(); + const descriptor = { + id: "claude-sdk", + displayName: "Claude", + getInstallState: () => ({ kind: "ready", source: "managed" }) as const, + subscribeInstallState: jest.fn(() => () => {}), + openInstallUI: jest.fn(), + createBackendProcess: jest.fn(() => procHandle.proc), + } as unknown as BackendDescriptor; + return { descriptor, procHandle }; +} + +describe("AgentModelPreloader.takeWarm", () => { + it("retains the probe subprocess after a successful preload and hands it to the manager", async () => { + const { descriptor, procHandle } = buildDescriptor(() => makeMockProc()); + const preloader = new AgentModelPreloader(buildApp(), buildPlugin(), () => descriptor); + + await preloader.preload("claude-sdk"); + + // Probe started but was NOT shut down — that's the whole point. + expect(procHandle.start).toHaveBeenCalledTimes(1); + expect(procHandle.shutdown).not.toHaveBeenCalled(); + + // State cache populated so the picker can read it before the manager + // takes ownership. + const cached = preloader.getCachedBackendState("claude-sdk"); + expect(cached?.model?.current.baseModelId).toBe("claude-sonnet"); + + // First takeWarm yields the running proc + probe sessionId + state. + const warm = preloader.takeWarm("claude-sdk"); + expect(warm).not.toBeNull(); + expect(warm?.proc).toBe(procHandle.proc); + expect(warm?.probeSessionId).toBe("probe-1"); + expect(warm?.state.model?.current.baseModelId).toBe("claude-sonnet"); + + // Single-shot — second call returns null so the manager can never + // hand out the same warm proc to two sessions. + expect(preloader.takeWarm("claude-sdk")).toBeNull(); + + // The cached state survives consumption so the picker keeps rendering + // until the live session pushes its own state via `setCached`. + expect(preloader.getCachedBackendState("claude-sdk")).not.toBeNull(); + }); + + it("shuts down a still-warm proc on dispose so the subprocess does not leak", async () => { + const { descriptor, procHandle } = buildDescriptor(() => makeMockProc()); + const preloader = new AgentModelPreloader(buildApp(), buildPlugin(), () => descriptor); + + await preloader.preload("claude-sdk"); + expect(procHandle.shutdown).not.toHaveBeenCalled(); + + preloader.shutdown(); + // shutdown is async-fire-and-forget; wait a tick. + await Promise.resolve(); + + expect(procHandle.shutdown).toHaveBeenCalledTimes(1); + expect(preloader.takeWarm("claude-sdk")).toBeNull(); + }); + + it("drops the warm entry when the probe subprocess exits before adoption", async () => { + const { descriptor, procHandle } = buildDescriptor(() => makeMockProc()); + const preloader = new AgentModelPreloader(buildApp(), buildPlugin(), () => descriptor); + + await preloader.preload("claude-sdk"); + expect(preloader.takeWarm("claude-sdk")).not.toBeNull(); + // Re-prime — populate a second warm entry to exercise the onExit path. + await preloader.preload("claude-sdk"); + procHandle.emitExit(); + + expect(preloader.takeWarm("claude-sdk")).toBeNull(); + }); + + it("shuts down the probe proc when the agent reports no usable state", async () => { + const { descriptor, procHandle } = buildDescriptor(() => + makeMockProc({ newSessionState: { model: null, mode: null } }) + ); + const preloader = new AgentModelPreloader(buildApp(), buildPlugin(), () => descriptor); + + await preloader.preload("claude-sdk"); + + // No usable catalog → preloader discards the proc so we don't keep a + // useless subprocess around. + expect(procHandle.shutdown).toHaveBeenCalledTimes(1); + expect(preloader.takeWarm("claude-sdk")).toBeNull(); + }); + + it("shuts down the probe proc when newSession throws", async () => { + const { descriptor, procHandle } = buildDescriptor(() => { + const handle = makeMockProc(); + handle.newSession.mockRejectedValueOnce(new Error("agent unreachable")); + return handle; + }); + const preloader = new AgentModelPreloader(buildApp(), buildPlugin(), () => descriptor); + + await preloader.preload("claude-sdk"); + + expect(procHandle.shutdown).toHaveBeenCalledTimes(1); + expect(preloader.takeWarm("claude-sdk")).toBeNull(); + }); + + it("clearCached shuts down and drops a still-warm proc", async () => { + const { descriptor, procHandle } = buildDescriptor(() => makeMockProc()); + const preloader = new AgentModelPreloader(buildApp(), buildPlugin(), () => descriptor); + + await preloader.preload("claude-sdk"); + expect(preloader.getCachedBackendState("claude-sdk")).not.toBeNull(); + + preloader.clearCached("claude-sdk"); + await Promise.resolve(); + + expect(procHandle.shutdown).toHaveBeenCalledTimes(1); + expect(preloader.getCachedBackendState("claude-sdk")).toBeNull(); + expect(preloader.takeWarm("claude-sdk")).toBeNull(); + }); +}); diff --git a/src/agentMode/session/AgentModelPreloader.ts b/src/agentMode/session/AgentModelPreloader.ts index 34793a66..1d38b5fb 100644 --- a/src/agentMode/session/AgentModelPreloader.ts +++ b/src/agentMode/session/AgentModelPreloader.ts @@ -4,23 +4,58 @@ import { getSettings } from "@/settings/model"; import { App, FileSystemAdapter, Platform } from "obsidian"; import { MethodUnsupportedError } from "./errors"; import { backendStateSignature } from "./translateBackendState"; -import type { BackendDescriptor, BackendId, BackendProcess, BackendState } from "./types"; +import type { + BackendDescriptor, + BackendId, + BackendProcess, + BackendState, + SessionId, +} from "./types"; /** - * Plugin-lifetime cache of per-backend session state. Backends expose - * `BackendState` only as a side-effect of session creation / resume / load, - * so without this preload the picker would show no entries for non-active - * backends and would blink empty during the round-trip on a fresh session. + * Warm result of a successful preload — the still-running backend process, + * the probe session id it owns, and the state snapshot it reported. Handed + * off via `takeWarm()` so the manager can skip a fresh subprocess spawn + + * `initialize` handshake on the first chat open, and (when the user opens + * a fresh chat on that backend) adopt the probe session as the user's + * session instead of paying another `newSession` round-trip. + */ +export interface WarmBackend { + proc: BackendProcess; + probeSessionId: SessionId; + state: BackendState; +} + +/** + * Plugin-lifetime cache of per-backend session state and the running probe + * subprocess that produced it. Backends expose `BackendState` only as a + * side-effect of session creation / resume / load, so without this preload + * the picker would show no entries for non-active backends and would blink + * empty during the round-trip on a fresh session. * * Probes once per backend at startup: prefer resume of a persisted probe * sessionId, fall back to load, then to new (and persist the new id so the * next reload can reuse it — keeps the agent-side session store at one stale * entry per machine instead of growing with each reload). + * + * The probe subprocess is **kept warm** until the manager consumes it via + * `takeWarm(backendId)`. That removes the warm subprocess spawn from the + * critical path of the first chat-open: instead of preload booting a + * subprocess just to read its catalog and immediately shutting it down, + * the same subprocess becomes the manager's backend process on first use. */ export class AgentModelPreloader { + private readonly warm = new Map(); + // State-only cache for backends whose warm entry has already been + // consumed (or that pushed updates after consumption via `setCached`). + // The active session's `attachModelCacheSync` writes here. private readonly cache = new Map(); private readonly inflight = new Map>(); private readonly listeners = new Set<() => void>(); + // Per-warm-entry exit-listener teardowns. Wired when the warm entry is + // recorded so we can clear it if the probe subprocess dies before the + // manager takes ownership. + private readonly warmExitUnsubs = new Map void>(); private disposed = false; constructor( @@ -30,26 +65,61 @@ export class AgentModelPreloader { ) {} getCachedBackendState(backendId: BackendId): BackendState | null { - return this.cache.get(backendId) ?? null; + return this.warm.get(backendId)?.state ?? this.cache.get(backendId) ?? null; } /** * Replace the cached entry for `backendId`. No-op when the signature - * is unchanged, to avoid spurious picker rebuilds. + * is unchanged, to avoid spurious picker rebuilds. Live sessions push + * here via `attachModelCacheSync` after the warm entry is consumed. */ setCached(backendId: BackendId, state: BackendState): void { if (this.disposed) return; - const prev = this.cache.get(backendId) ?? null; + const prev = this.getCachedBackendState(backendId); if (backendStateSignature(prev) === backendStateSignature(state)) return; this.cache.set(backendId, state); this.notify(); } - /** Remove the cached entry for `backendId` after its backend is restarted. */ + /** + * Remove all cached state for `backendId` after its backend is restarted. + * Drops the warm subprocess if it hasn't been taken yet so a fresh probe + * runs on the next `preload(backendId)` call. + */ clearCached(backendId: BackendId): void { if (this.disposed) return; - if (!this.cache.delete(backendId)) return; - this.notify(); + let changed = this.cache.delete(backendId); + const warm = this.warm.get(backendId); + if (warm) { + this.warm.delete(backendId); + this.warmExitUnsubs.get(backendId)?.(); + this.warmExitUnsubs.delete(backendId); + // Best-effort shutdown of the abandoned warm proc. + warm.proc.shutdown().catch((e) => { + logWarn(`[AgentMode] preload clearCached: shutdown of warm ${backendId} failed`, e); + }); + changed = true; + } + if (changed) this.notify(); + } + + /** + * Hand the warm backend (process + probe session id + state snapshot) to + * the manager. Single-shot: removes the entry so subsequent callers see + * `null` and the manager owns lifetime of the process from here on. + */ + takeWarm(backendId: BackendId): WarmBackend | null { + const entry = this.warm.get(backendId); + if (!entry) return null; + this.warm.delete(backendId); + this.warmExitUnsubs.get(backendId)?.(); + this.warmExitUnsubs.delete(backendId); + // Keep the last-known state in `cache` so the picker still reads it + // until the live session syncs its own state back in. Without this the + // picker would briefly see `null` between `takeWarm` and the first + // `attachModelCacheSync` write. + this.cache.set(backendId, entry.state); + return entry; } /** Best-effort probe; failures are logged and swallowed. Dedupes per backend. */ @@ -74,6 +144,14 @@ export class AgentModelPreloader { this.cache.clear(); this.inflight.clear(); this.listeners.clear(); + for (const [backendId, warm] of this.warm) { + this.warmExitUnsubs.get(backendId)?.(); + warm.proc.shutdown().catch((e) => { + logWarn(`[AgentMode] preload shutdown: warm ${backendId} shutdown failed`, e); + }); + } + this.warm.clear(); + this.warmExitUnsubs.clear(); } private notify(): void { @@ -106,26 +184,50 @@ export class AgentModelPreloader { descriptor, }); + let probe: { sessionId: SessionId; state: BackendState } | null = null; try { await proc.start?.(); const storedId = descriptor.getProbeSessionId?.(getSettings()); - const state = await this.fetchInitialState(proc, descriptor, backendId, storedId, cwd); - if (this.disposed) return; - if (state.model || state.mode) { - this.setCached(backendId, state); - logProbeResult(backendId, "session probe", state); - } else { - logInfo(`[AgentMode] preload ${backendId}: agent did not report any initial state`); - } + probe = await this.fetchInitialState(proc, descriptor, backendId, storedId, cwd); } catch (err) { logError(`[AgentMode] preload ${backendId} failed`, err); - } finally { + } + + if (this.disposed || !probe || (!probe.state.model && !probe.state.mode)) { + if (probe) { + logInfo(`[AgentMode] preload ${backendId}: agent did not report any initial state`); + } try { await proc.shutdown(); } catch (e) { logWarn(`[AgentMode] preload ${backendId}: shutdown failed`, e); } + return; } + + // Probe succeeded — retain the running subprocess as a warm entry so + // the first chat-open can adopt it instead of paying another spawn + + // initialize round-trip. + const warm: WarmBackend = { + proc, + probeSessionId: probe.sessionId, + state: probe.state, + }; + const exitUnsub = proc.onExit(() => { + if (this.disposed) return; + // Subprocess died before the manager claimed it. Drop the warm + // entry; next createSession will spawn a fresh one through the + // descriptor. + if (this.warm.get(backendId) === warm) { + this.warm.delete(backendId); + this.warmExitUnsubs.delete(backendId); + this.notify(); + } + }); + this.warm.set(backendId, warm); + this.warmExitUnsubs.set(backendId, exitUnsub); + logProbeResult(backendId, "session probe", probe.state); + this.notify(); } private async fetchInitialState( @@ -134,7 +236,7 @@ export class AgentModelPreloader { backendId: BackendId, storedId: string | undefined, cwd: string - ): Promise { + ): Promise<{ sessionId: SessionId; state: BackendState }> { type Strategy = { label: string; sessionId: string; @@ -156,10 +258,14 @@ export class AgentModelPreloader { for (const { label, sessionId, run } of strategies) { try { + // Register a no-op handler before the call so updates emitted + // during the call are demuxed against this sessionId rather than + // buffered indefinitely. The manager overrides this with the real + // handler when it adopts the session. proc.registerSessionHandler(sessionId, () => {}); const resp = await run(); logInfo(`[AgentMode] preload ${backendId}: ${label}`); - return resp.state; + return { sessionId: resp.sessionId, state: resp.state }; } catch (err) { if (!(err instanceof MethodUnsupportedError)) { logWarn(`[AgentMode] preload ${backendId}: ${label} failed (will fall back)`, err); @@ -177,7 +283,7 @@ export class AgentModelPreloader { logWarn(`[AgentMode] preload ${backendId}: persistProbeSessionId failed`, e); } } - return resp.state; + return { sessionId: resp.sessionId, state: resp.state }; } } diff --git a/src/agentMode/session/AgentSession.test.ts b/src/agentMode/session/AgentSession.test.ts index 8e766ea1..cb8cc066 100644 --- a/src/agentMode/session/AgentSession.test.ts +++ b/src/agentMode/session/AgentSession.test.ts @@ -490,15 +490,31 @@ describe("AgentSession.create (via start)", () => { expect(session.getState()?.model).toBeNull(); }); - it("attempts setModel when defaultModelId is set", async () => { + it("attempts setModel when defaultModelSelection is set", async () => { const mock = makeMockBackend(); - mock.newSession.mockResolvedValueOnce({ sessionId: "acp-1", state: emptyState() }); + const stateWithSonnet: BackendState = { + model: { + current: { baseModelId: "anthropic/sonnet", effort: null }, + availableModels: [ + { + baseModelId: "anthropic/sonnet", + name: "Claude Sonnet", + provider: "anthropic", + effortOptions: [], + }, + { baseModelId: "openai/gpt-5", name: "GPT-5", provider: "openai", effortOptions: [] }, + ], + }, + mode: null, + }; + mock.newSession.mockResolvedValueOnce({ sessionId: "acp-1", state: stateWithSonnet }); const session = AgentSession.start({ backend: mock.asBackend, cwd: "/vault", internalId: "internal-1", backendId: "opencode", - defaultModelId: "openai/gpt-5", + defaultModelSelection: { baseModelId: "openai/gpt-5", effort: null }, + getDescriptor: () => makeWireOnlyDescriptor(), }); await session.ready; expect(mock.setSessionModel).toHaveBeenCalledWith({ @@ -507,22 +523,276 @@ describe("AgentSession.create (via start)", () => { }); }); - it("survives a MethodUnsupportedError from default-model application", async () => { + it("seeds currentState with the persisted selection before notifying listeners", async () => { const mock = makeMockBackend(); - mock.newSession.mockResolvedValueOnce({ sessionId: "acp-1", state: emptyState() }); + const stateWithSonnet: BackendState = { + model: { + current: { baseModelId: "anthropic/sonnet", effort: null }, + availableModels: [ + { + baseModelId: "anthropic/sonnet", + name: "Claude Sonnet", + provider: "anthropic", + effortOptions: [], + }, + { baseModelId: "openai/gpt-5", name: "GPT-5", provider: "openai", effortOptions: [] }, + ], + }, + mode: null, + }; + mock.newSession.mockResolvedValueOnce({ sessionId: "acp-1", state: stateWithSonnet }); + // Block setSessionModel so the seed must survive on its own — without + // the optimistic seed the picker would see "anthropic/sonnet" first. + let resolveSetModel: ((s: BackendState) => void) | null = null; + mock.setSessionModel.mockImplementationOnce( + () => new Promise((resolve) => (resolveSetModel = resolve)) + ); + const session = AgentSession.start({ + backend: mock.asBackend, + cwd: "/vault", + internalId: "internal-1", + backendId: "opencode", + defaultModelSelection: { baseModelId: "openai/gpt-5", effort: null }, + getDescriptor: () => makeWireOnlyDescriptor(), + }); + + const observed: Array = []; + session.subscribe({ + onMessagesChanged: jest.fn(), + onStatusChanged: jest.fn(), + onModelChanged: () => observed.push(session.getState()?.model?.current.baseModelId), + }); + + // Wait for the first notifyModelChanged inside initialize. + await Promise.resolve(); + await Promise.resolve(); + expect(observed[0]).toBe("openai/gpt-5"); + + resolveSetModel!(stateWithSonnet); + await session.ready; + }); + + it("eagerly seeds currentState from initialCachedState before newSession resolves", async () => { + const mock = makeMockBackend(); + const cachedState: BackendState = { + model: { + current: { baseModelId: "kimi-2.6", effort: null }, + availableModels: [ + { baseModelId: "kimi-2.6", name: "Kimi 2.6", provider: "moon", effortOptions: [] }, + { + baseModelId: "big-pickle", + name: "Big Pickle", + provider: "moon", + effortOptions: [], + }, + ], + }, + mode: null, + }; + // Block newSession so we can observe the pre-initialize state. + let resolveNewSession: ((r: { sessionId: string; state: BackendState }) => void) | null = null; + mock.newSession.mockImplementationOnce( + () => + new Promise<{ sessionId: string; state: BackendState }>( + (resolve) => (resolveNewSession = resolve) + ) + ); + const session = AgentSession.start({ + backend: mock.asBackend, + cwd: "/vault", + internalId: "internal-1", + backendId: "opencode", + defaultModelSelection: { baseModelId: "big-pickle", effort: null }, + initialCachedState: cachedState, + getDescriptor: () => makeWireOnlyDescriptor(), + }); + + // Before newSession resolves, getState reflects the eager seed + // (current = big-pickle) rather than the cached current (kimi-2.6). + expect(session.getState()?.model?.current.baseModelId).toBe("big-pickle"); + + resolveNewSession!({ sessionId: "acp-1", state: cachedState }); + await session.ready; + }); + + it("does not seed effort into currentState (effort stays as backend reported)", async () => { + // Regression: previously, the optimistic seed wrote both `baseModelId` + // and `effort` into currentState. For descriptor-style backends where + // effort lives outside the wire id, `applyInitialSessionConfig` would + // see the seeded effort match the persisted effort and skip the real + // `setConfigOption`, silently dropping the user's persisted effort. + const mock = makeMockBackend(); + const backendState: BackendState = { + model: { + current: { baseModelId: "anthropic/sonnet", effort: "low" }, + availableModels: [ + { + baseModelId: "anthropic/sonnet", + name: "Claude Sonnet", + provider: "anthropic", + effortOptions: [ + { value: "low", label: "Low" }, + { value: "high", label: "High" }, + ], + }, + ], + }, + mode: null, + }; + mock.newSession.mockResolvedValueOnce({ sessionId: "acp-1", state: backendState }); + const session = AgentSession.start({ + backend: mock.asBackend, + cwd: "/vault", + internalId: "internal-1", + backendId: "claude", + defaultModelSelection: { baseModelId: "anthropic/sonnet", effort: "high" }, + // Descriptor whose wire encoding ignores effort (Claude-style). + getDescriptor: () => makeDescriptorWireWithoutEffort(), + }); + await session.ready; + // baseModelId matches → no setModel call needed. + expect(mock.setSessionModel).not.toHaveBeenCalled(); + // Effort is what the backend reported, NOT the persisted "high" — so + // `applyInitialSessionConfig` will see a mismatch and call setConfigOption. + expect(session.getState()?.model?.current.effort).toBe("low"); + }); + + it("reverts the seeded selection when setModel fails", async () => { + const mock = makeMockBackend(); + const stateWithSonnet: BackendState = { + model: { + current: { baseModelId: "anthropic/sonnet", effort: null }, + availableModels: [ + { + baseModelId: "anthropic/sonnet", + name: "Claude Sonnet", + provider: "anthropic", + effortOptions: [], + }, + { baseModelId: "openai/gpt-5", name: "GPT-5", provider: "openai", effortOptions: [] }, + ], + }, + mode: null, + }; + mock.newSession.mockResolvedValueOnce({ sessionId: "acp-1", state: stateWithSonnet }); mock.setSessionModel.mockRejectedValueOnce(new MethodUnsupportedError("session/set_model")); const session = AgentSession.start({ backend: mock.asBackend, cwd: "/vault", internalId: "internal-1", backendId: "opencode", - defaultModelId: "openai/gpt-5", + defaultModelSelection: { baseModelId: "openai/gpt-5", effort: null }, + getDescriptor: () => makeWireOnlyDescriptor(), }); await session.ready; expect(session.getStatus()).toBe("idle"); + // Seed reverted to whatever the backend actually reported. + expect(session.getState()?.model?.current.baseModelId).toBe("anthropic/sonnet"); }); }); +/** Minimal wire-only descriptor for tests that exercise seed/setModel. */ +function makeWireOnlyDescriptor(): BackendDescriptor { + return { + wire: { + encode: (selection: { baseModelId: string; effort: string | null }) => + selection.effort ? `${selection.baseModelId}/${selection.effort}` : selection.baseModelId, + decode: (wireId: string) => ({ + selection: { baseModelId: wireId, effort: null }, + provider: null, + }), + }, + } as unknown as BackendDescriptor; +} + +describe("AgentSession warm-adoption ready gating", () => { + it("ready stays pending until setModel resolves so sendPrompt can't fire on the probe model", async () => { + const mock = makeMockBackend(); + const probeState: BackendState = { + model: { + current: { baseModelId: "anthropic/sonnet", effort: null }, + availableModels: [ + { + baseModelId: "anthropic/sonnet", + name: "Sonnet", + provider: "anthropic", + effortOptions: [], + }, + { baseModelId: "openai/gpt-5", name: "GPT-5", provider: "openai", effortOptions: [] }, + ], + }, + mode: null, + }; + // Block setSessionModel so we can observe `ready` remaining pending. + let resolveSetModel!: (s: BackendState) => void; + mock.setSessionModel.mockImplementationOnce( + () => new Promise((resolve) => (resolveSetModel = resolve)) + ); + + const session = new AgentSession({ + backend: mock.asBackend, + backendSessionId: "probe-1", + internalId: "internal-1", + backendId: "opencode", + initialState: probeState, + defaultModelSelection: { baseModelId: "openai/gpt-5", effort: null }, + getDescriptor: () => makeWireOnlyDescriptor(), + }); + + let readyResolved = false; + void session.ready.then(() => { + readyResolved = true; + }); + await new Promise((r) => window.setTimeout(r, 0)); + expect(readyResolved).toBe(false); + expect(mock.setSessionModel).toHaveBeenCalledWith({ + sessionId: "probe-1", + modelId: "openai/gpt-5", + }); + + resolveSetModel({ + model: { + ...probeState.model!, + current: { baseModelId: "openai/gpt-5", effort: null }, + }, + mode: null, + }); + await session.ready; + expect(readyResolved).toBe(true); + expect(session.getState()?.model?.current.baseModelId).toBe("openai/gpt-5"); + }); + + it("ready resolves immediately when no default selection is supplied", async () => { + const mock = makeMockBackend(); + const session = new AgentSession({ + backend: mock.asBackend, + backendSessionId: "probe-1", + internalId: "internal-1", + backendId: "opencode", + initialState: emptyState(), + }); + await session.ready; + expect(mock.setSessionModel).not.toHaveBeenCalled(); + }); +}); + +/** + * Descriptor whose wire encoding ignores effort (Claude SDK-style). Used + * to exercise the "effort lives out-of-band" branch where `setModel` + * would be a no-op and effort is applied via `applyInitialSessionConfig`. + */ +function makeDescriptorWireWithoutEffort(): BackendDescriptor { + return { + wire: { + encode: (selection: { baseModelId: string; effort: string | null }) => selection.baseModelId, + decode: (wireId: string) => ({ + selection: { baseModelId: wireId, effort: null }, + provider: null, + }), + }, + } as unknown as BackendDescriptor; +} + describe("AgentSession.setModel", () => { it("calls backend.setSessionModel and replaces the cached state on success", async () => { const mock = makeMockBackend(); diff --git a/src/agentMode/session/AgentSession.ts b/src/agentMode/session/AgentSession.ts index bae6d31c..a07972ab 100644 --- a/src/agentMode/session/AgentSession.ts +++ b/src/agentMode/session/AgentSession.ts @@ -9,6 +9,7 @@ import { BackendProcess, BackendState, CurrentPlan, + ModelSelection, NewAgentChatMessage, PERMISSION_ALLOW_KINDS, PERMISSION_REJECT_KINDS, @@ -43,6 +44,37 @@ const MAX_TOOL_OUTPUT_TEXT_CHARS = 12_000; // behavior on idle subscription ticks. const EMPTY_PERMISSIONS: PermissionPrompt[] = []; +/** + * Optimistically swap `state.model.current.baseModelId` for the persisted + * user selection so the picker's first paint matches what the user picked. + * + * Only `baseModelId` is seeded — `effort` is left as the backend reported. + * For descriptor-style backends (Claude) effort lives out-of-band and is + * applied via `applyInitialSessionConfig` → `setConfigOption`; seeding it + * here would make that step see a matching value and silently skip the + * real config write. For wire-effort backends (opencode-style) the + * subsequent `setModel` call carries the user's effort through the + * encoded id, and the backend's response replaces this seed entirely. + * + * Returns the input reference unchanged when no seed is possible. + */ +function seedSelectionIntoState( + state: BackendState | null, + selection: ModelSelection | undefined +): BackendState | null { + if (!state || !state.model || !selection) return state; + const entry = state.model.availableModels.find((m) => m.baseModelId === selection.baseModelId); + if (!entry) return state; + if (state.model.current.baseModelId === selection.baseModelId) return state; + return { + ...state, + model: { + ...state.model, + current: { ...state.model.current, baseModelId: selection.baseModelId }, + }, + }; +} + export type AgentSessionStatus = | "starting" | "idle" @@ -94,7 +126,21 @@ export interface AgentSessionStartOptions { cwd: string; internalId: string; backendId: BackendId; - defaultModelId?: string; + /** + * Persisted user preference to apply after the backend's initial session + * state. The session seeds it optimistically so the first picker paint + * shows the user's pick, then confirms with the backend via setModel + * (and, for descriptor-style backends, setConfigOption — handled by the + * manager's `applyInitialSessionConfig` hook). + */ + defaultModelSelection?: ModelSelection; + /** + * Snapshot from the preloader cache used to seed `currentState` + * synchronously so the picker doesn't flash the previous session's + * selection during the `backend.newSession` round-trip. Replaced by the + * agent's fresh `newSession` response inside `initialize()`. + */ + initialCachedState?: BackendState | null; /** * Optional descriptor accessor. The session uses it to resolve mode mappings * without coupling to specific backends. Manager-supplied; tests omit it. @@ -112,6 +158,13 @@ export interface AgentSessionStateOptions { internalId: string; backendId: BackendId; initialState?: BackendState | null; + /** + * Optional persisted user preference applied to the warm/adopted session. + * Same role as `AgentSessionStartOptions.defaultModelSelection` — seeded + * optimistically on construction; setModel runs async and reverts the + * seed on failure. + */ + defaultModelSelection?: ModelSelection; cwd?: string | null; getDescriptor?: () => BackendDescriptor | undefined; } @@ -219,19 +272,30 @@ export class AgentSession { this.getDescriptor = opts.getDescriptor ?? null; if ("backendSessionId" in opts) { this.backendSessionId = opts.backendSessionId; - this.currentState = opts.initialState ?? null; + const originalState = opts.initialState ?? null; + this.currentState = seedSelectionIntoState(originalState, opts.defaultModelSelection); this.unregisterSessionHandler = this.backend.registerSessionHandler( opts.backendSessionId, (event) => this.handleSessionEvent(event) ); - // backendSessionId is non-null → getStatus() yields "idle" without - // any further bookkeeping. Sync the cache so the first - // `recomputeStatusIfChanged` call doesn't fire a spurious - // "starting → idle" transition that no one observed. + // Sync the status cache so the first recomputeStatusIfChanged doesn't + // fire a spurious "starting → idle" transition that no one observed. this.cachedStatus = this.getStatus(); - this.ready = Promise.resolve(); + // Gate `ready` on the model confirmation round-trip so `sendPrompt` + // can't fire on the probe's model before the user's persisted + // selection is applied to the backend. + this.ready = + opts.defaultModelSelection && originalState + ? this.confirmSeededSelection(opts.defaultModelSelection, originalState) + : Promise.resolve(); } else { - // Default cachedStatus ("starting") already matches getStatus(). + // Eagerly seed from the preloader cache so the picker doesn't fall + // back to the prior session's `current` while `backend.newSession` is + // in flight. `initialize()` replaces this with the agent's response. + this.currentState = seedSelectionIntoState( + opts.initialCachedState ?? null, + opts.defaultModelSelection + ); this.ready = this.initialize(opts); } } @@ -252,7 +316,7 @@ export class AgentSession { } private async initialize(opts: AgentSessionStartOptions): Promise { - const { backend, cwd, defaultModelId } = opts; + const { backend, cwd, defaultModelSelection } = opts; try { const resp = await backend.newSession({ cwd, @@ -266,7 +330,7 @@ export class AgentSession { : "agent did not report model state"; logInfo(`[AgentMode] session ${resp.sessionId} ${modelLog}`); this.backendSessionId = resp.sessionId; - this.currentState = resp.state; + this.currentState = seedSelectionIntoState(resp.state, defaultModelSelection); this.unregisterSessionHandler = this.backend.registerSessionHandler(resp.sessionId, (event) => this.handleSessionEvent(event) ); @@ -279,14 +343,8 @@ export class AgentSession { this.recomputeStatusIfChanged(); this.notifyModelChanged(); - // Apply sticky preference. Best-effort — failures leave the session - // usable with whatever the agent picked by default. - if (defaultModelId) { - try { - await this.setModel(defaultModelId); - } catch (e) { - logWarn(`[AgentMode] could not apply default model ${defaultModelId}`, e); - } + if (defaultModelSelection) { + await this.confirmSeededSelection(defaultModelSelection, resp.state); } } catch (err) { if (this.disposed) return; @@ -326,6 +384,42 @@ export class AgentSession { this.notifyModelChanged(); } + /** + * Apply the persisted user selection to the backend via `setModel`. + * Runs whenever `defaultModelSelection` is supplied — `setModel` is + * idempotent, and the wire-encoding short-circuit below makes it a + * pure no-op when the selection already matches the backend's report. + * + * On success the backend's response replaces `currentState`. On + * failure any optimistic baseModelId seed is reverted to `originalState` + * so the picker drops back to whatever the backend actually has. + * + * Skips the round-trip when the encoded form matches. For descriptor- + * style backends (Claude SDK) where effort lives outside the wire id, + * an effort-only change encodes the same and `setModel` is correctly + * skipped here — `applyInitialSessionConfig` dispatches the effort + * via `setConfigOption` instead. + */ + private async confirmSeededSelection( + selection: ModelSelection, + originalState: BackendState + ): Promise { + const descriptor = this.getDescriptor?.(); + if (!descriptor) return; + const encoded = descriptor.wire.encode(selection); + const originalEncoded = originalState.model + ? descriptor.wire.encode(originalState.model.current) + : null; + if (encoded === originalEncoded) return; + try { + await this.setModel(encoded); + } catch (e) { + logWarn(`[AgentMode] could not apply default model ${encoded}; reverting seed`, e); + this.currentState = originalState; + this.notifyModelChanged(); + } + } + /** * Set a session configuration option (e.g. effort). Reuses * `notifyModelChanged` because the picker treats model and configOption diff --git a/src/agentMode/session/AgentSessionManager.test.ts b/src/agentMode/session/AgentSessionManager.test.ts index 11a5ac15..6f40c859 100644 --- a/src/agentMode/session/AgentSessionManager.test.ts +++ b/src/agentMode/session/AgentSessionManager.test.ts @@ -38,6 +38,9 @@ function makeMockBackendProcess() { }, isRunning: () => mockBackendIsRunning, shutdown: mockBackendShutdown, + // Stub the session-event surface so warm-adoption tests can construct + // a real `AgentSession` via the state-options branch without throwing. + registerSessionHandler: jest.fn(() => () => {}), }; } @@ -150,6 +153,7 @@ function buildManager(): AgentSessionManager { shutdown: jest.fn(), setCached: jest.fn(), clearCached: jest.fn(), + takeWarm: jest.fn(() => null), }; return new AgentSessionManager( buildApp(), @@ -223,6 +227,7 @@ describe("AgentSessionManager.createSession", () => { clearCached: jest.fn((id: string) => { cache.delete(id); }), + takeWarm: jest.fn(() => null), }; const descriptor = buildDescriptor(); const mgr = new AgentSessionManager( @@ -307,6 +312,121 @@ describe("AgentSessionManager.createSession", () => { }); }); +describe("AgentSessionManager warm-backend adoption", () => { + it("reuses the preloader's warm proc instead of spawning a fresh one", async () => { + // A warm proc is just another mock backend — the manager doesn't care + // it came from the preloader, only that it skips start() / newSession. + const warmProc = makeMockBackendProcess(); + const descriptor = buildDescriptor(); + const probeState = { + model: { + current: { baseModelId: "anthropic/sonnet", effort: null }, + availableModels: [ + { baseModelId: "anthropic/sonnet", name: "Sonnet", provider: null, effortOptions: [] }, + ], + }, + mode: null, + }; + const takeWarmMock = jest + .fn() + .mockReturnValueOnce({ + proc: warmProc, + probeSessionId: "probe-1", + state: probeState, + }) + .mockReturnValue(null); + const modelPreloader = { + getCachedBackendState: jest.fn(() => probeState), + preload: jest.fn(async () => undefined), + subscribe: jest.fn(() => () => {}), + shutdown: jest.fn(), + setCached: jest.fn(), + clearCached: jest.fn(), + takeWarm: takeWarmMock, + }; + const mgr = new AgentSessionManager( + buildApp(), + buildPlugin() as unknown as ConstructorParameters[1], + { + permissionPrompter: jest.fn(), + resolveDescriptor: (id) => (id === descriptor.id ? descriptor : undefined), + modelPreloader: modelPreloader as unknown as ConstructorParameters< + typeof AgentSessionManager + >[2]["modelPreloader"], + } + ); + + await mgr.createSession(); + + // start() on the warm proc must not be re-invoked — preload already did it. + expect(mockBackendStart).not.toHaveBeenCalled(); + // The session should NOT have gone through AgentSession.start (that path + // is for newSession-driven creates). Warm adoption uses the constructor + // directly. + expect(sessionCreateSpy).not.toHaveBeenCalled(); + // No fresh backend was created either — descriptor.createBackendProcess + // is the spawn primitive. + expect(descriptor.createBackendProcess).not.toHaveBeenCalled(); + // Active session is the warm one. + expect(mgr.getActiveSession()).not.toBeNull(); + expect(mgr.getActiveSession()?.backendId).toBe("opencode"); + }); + + it("falls back to a fresh spawn when no warm entry is available", async () => { + // takeWarm returning null is the second-and-later session path. + const mgr = buildManager(); + await mgr.createSession(); + + // No warm entry → start() runs on the freshly-created backend. + expect(mockBackendStart).toHaveBeenCalledTimes(1); + // AgentSession.start was used, not the construct-with-state branch. + expect(sessionCreateSpy).toHaveBeenCalledTimes(1); + }); +}); + +describe("AgentSessionManager preload status", () => { + it("isPreloadReady reports per-backend; a missing entry is treated as ready", () => { + const mgr = buildManager(); + // Unregistered backend → ready (so the chat doesn't stall waiting on a + // preload that will never run). + expect(mgr.isPreloadReady("opencode")).toBe(true); + expect(mgr.getPreloadStatus("opencode")).toBe("absent"); + }); + + it("transitions pending → ready on promise resolution and ready → ready stays ready", async () => { + const mgr = buildManager(); + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + mgr.registerPreload("opencode", promise); + expect(mgr.getPreloadStatus("opencode")).toBe("pending"); + expect(mgr.isPreloadReady("opencode")).toBe(false); + resolve(); + await promise; + await Promise.resolve(); + expect(mgr.getPreloadStatus("opencode")).toBe("ready"); + expect(mgr.isPreloadReady("opencode")).toBe(true); + }); + + it("transitions pending → error on promise rejection but still reports ready (chat unblocks)", async () => { + const mgr = buildManager(); + const rejection = new Error("preload failed"); + mgr.registerPreload( + "opencode", + Promise.reject(rejection).catch((e) => { + throw e; + }) + ); + // Microtask flush. + await new Promise((r) => window.setTimeout(r, 0)); + expect(mgr.getPreloadStatus("opencode")).toBe("error"); + // The chat treats error as "ready enough" so the user isn't stuck on + // a perpetual spinner; the picker shows the failure row. + expect(mgr.isPreloadReady("opencode")).toBe(true); + }); +}); + describe("AgentSessionManager.getOrCreateActiveSession", () => { it("dedupes concurrent auto-spawn callers into a single session", async () => { const mgr = buildManager(); @@ -668,6 +788,7 @@ describe("AgentSessionManager.applySelection", () => { shutdown: jest.fn(), setCached: jest.fn(), clearCached: jest.fn(), + takeWarm: jest.fn(() => null), }; const mgr = new AgentSessionManager( buildApp(), @@ -758,6 +879,7 @@ describe("AgentSessionManager.applySelection", () => { shutdown: jest.fn(), setCached: jest.fn(), clearCached: jest.fn(), + takeWarm: jest.fn(() => null), }; const mgr = new AgentSessionManager( buildApp(), diff --git a/src/agentMode/session/AgentSessionManager.ts b/src/agentMode/session/AgentSessionManager.ts index f2e7edf5..8345a2a7 100644 --- a/src/agentMode/session/AgentSessionManager.ts +++ b/src/agentMode/session/AgentSessionManager.ts @@ -9,7 +9,7 @@ import { App, FileSystemAdapter, Notice, Platform, TFile } from "obsidian"; import { v4 as uuidv4 } from "uuid"; import { AgentSession, ATTENTION_TRIGGER_STATUSES } from "./AgentSession"; import type { AgentChatPersistenceManager } from "./AgentChatPersistenceManager"; -import type { AgentModelPreloader } from "./AgentModelPreloader"; +import type { AgentModelPreloader, WarmBackend } from "./AgentModelPreloader"; import { MethodUnsupportedError } from "./errors"; import { resolveMcpServers } from "./mcpResolver"; import type { @@ -74,12 +74,13 @@ export class AgentSessionManager { private readonly restartingBackends = new Set(); private readonly preloader: AgentModelPreloader; /** - * Resolves once the plugin-load preload phase has settled (regardless of - * per-backend success). Lets the chat UI gate its first render so the - * picker never flashes empty before the cache populates. + * Per-backend preload status. The chat UI gates its first render on the + * *active* backend's status; the model picker shows a "Loading…" or + * "Failed to load" placeholder for any backend that hasn't reached + * `"ready"` yet. A missing entry is treated as `"absent"` (backend not + * installed / never queued) which is render-safe. */ - private preloadPromise: Promise = Promise.resolve(); - private preloadReady = true; + private readonly preloadStatus = new Map(); // Per-session bookkeeping, all keyed by `internalId`. Mixes persistence // bookkeeping with subscription teardowns — the unifying property is // "must be cleaned up when the session is detached": @@ -201,8 +202,12 @@ export class AgentSessionManager { this.notify(); let backend: BackendProcess; + let warm: WarmBackend | null = null; try { - backend = await this.ensureBackend(resolvedId, descriptor); + // ensureBackend tries the preloader's warm probe before paying a + // fresh spawn. When `warm` is non-null the probe session is + // available for adoption below — skips a second `newSession`. + ({ proc: backend, warm } = await this.ensureBackend(resolvedId, descriptor)); } catch (err) { this.lastError = err2String(err); this.finishPendingCreate(); @@ -214,16 +219,39 @@ export class AgentSessionManager { throw new Error("AgentSessionManager was shut down during session creation"); } - const seedSelection = this.getDefaultSelection(resolvedId); - const defaultModelId = seedSelection ? descriptor.wire.encode(seedSelection) : undefined; - const session = AgentSession.start({ - backend, - cwd: vaultBasePath, - internalId: uuidv4(), - backendId: resolvedId, - defaultModelId, - getDescriptor: () => this.opts.resolveDescriptor(resolvedId), - }); + const seedSelection = this.getDefaultSelection(resolvedId) ?? undefined; + + let session: AgentSession; + if (warm) { + // The probe session that preload created is reused as the user's + // session: skips an extra `newSession` round-trip on the warm proc. + // The probe ran with no user preferences, so `defaultModelSelection` + // is what makes the session apply the persisted preference (with the + // optimistic seed so the picker never flashes the probe's default). + session = new AgentSession({ + backend, + backendSessionId: warm.probeSessionId, + internalId: uuidv4(), + backendId: resolvedId, + initialState: warm.state, + defaultModelSelection: seedSelection, + cwd: vaultBasePath, + getDescriptor: () => this.opts.resolveDescriptor(resolvedId), + }); + logInfo( + `[AgentMode] session adopted warm probe (internal=${session.internalId} backend-id=${warm.probeSessionId} backend=${resolvedId})` + ); + } else { + session = AgentSession.start({ + backend, + cwd: vaultBasePath, + internalId: uuidv4(), + backendId: resolvedId, + defaultModelSelection: seedSelection, + initialCachedState: this.preloader.getCachedBackendState(resolvedId), + getDescriptor: () => this.opts.resolveDescriptor(resolvedId), + }); + } this.sessions.set(session.internalId, session); this.chatUIStates.set(session.internalId, new AgentChatUIState(session)); this.activeSessionId = session.internalId; @@ -236,7 +264,8 @@ export class AgentSessionManager { // (claude-code's effort, future config-option preferences) and clear the // "starting" pill. On failure, capture into `lastError` so the status // surface and retry handler can react. The session itself transitions to - // status "error" inside its own `initialize`. + // status "error" inside its own `initialize`. For warm-adopted sessions + // `ready` is already resolved, so the chain runs on the next microtask. void session.ready .then(async () => { if (descriptor.applyInitialSessionConfig) { @@ -419,29 +448,52 @@ export class AgentSessionManager { } /** - * Register the aggregate preload promise — typically the `Promise.allSettled` - * of every backend's `preloadModels` call from plugin bring-up. While it - * pends, `isPreloadReady()` returns `false`; the chat UI uses that flag to - * render a "Loading…" placeholder instead of an empty picker. + * Register a backend's plugin-load preload promise. While the promise is + * pending, `getPreloadStatus(backendId)` returns `"pending"`; on settle it + * transitions to `"ready"` (the preload itself is best-effort, so even + * rejected promises flip to `"ready"` here — the picker reads the cached + * state alongside this status and renders empty-but-error vs loading). + * On programmatic error (rejected promise) the entry becomes `"error"` so + * the picker can offer a retry affordance. */ - setPreloadPromise(promise: Promise): void { - this.preloadReady = false; - this.preloadPromise = promise; - void promise.then(() => { - this.preloadReady = true; - this.notify(); - }); + registerPreload(backendId: BackendId, promise: Promise): void { + this.preloadStatus.set(backendId, "pending"); this.notify(); + promise.then( + () => { + if (this.disposed) return; + this.preloadStatus.set(backendId, "ready"); + this.notify(); + }, + () => { + if (this.disposed) return; + this.preloadStatus.set(backendId, "error"); + this.notify(); + } + ); } - /** Resolves once `setPreloadPromise`'s promise settles. */ - whenPreloadReady(): Promise { - return this.preloadPromise; + /** + * Synchronous check, suitable for React render gates. Defaults to the + * active backend; pass `backendId` to query a specific backend (the + * picker reads each backend's status). Backends that were never queued + * (not installed / disabled) are treated as ready so the chat UI doesn't + * stall waiting on a preload that will never run. + */ + isPreloadReady(backendId?: BackendId): boolean { + const id = backendId ?? getSettings().agentMode?.activeBackend; + if (!id) return true; + const status = this.preloadStatus.get(id); + return status === undefined || status === "ready" || status === "error"; } - /** Synchronous check, suitable for React render gates. */ - isPreloadReady(): boolean { - return this.preloadReady; + /** + * Read a backend's preload status for the picker. `"absent"` indicates + * the backend was never queued (uninstalled or disabled) — the picker + * uses today's behavior in that case (omit the section silently). + */ + getPreloadStatus(backendId: BackendId): "pending" | "ready" | "error" | "absent" { + return this.preloadStatus.get(backendId) ?? "absent"; } /** @@ -659,6 +711,7 @@ export class AgentSessionManager { this.starting.clear(); this.startingBackendId = null; this.listeners.clear(); + this.preloadStatus.clear(); this.preloader.shutdown(); } @@ -728,7 +781,10 @@ export class AgentSessionManager { let backend: BackendProcess; try { - backend = await this.ensureBackend(backendId, descriptor); + // Discard the optional warm probe-session info: we're rehydrating a + // specific saved session, not opening a fresh chat. The probe + // session sits unused on the proc until shutdown — harmless. + ({ proc: backend } = await this.ensureBackend(backendId, descriptor)); } catch (err) { this.lastError = err2String(err); this.finishPendingCreate(); @@ -954,14 +1010,40 @@ export class AgentSessionManager { this.getSessionState(session.internalId).attentionUnsub = unsubscribe; } + /** + * Obtain a running backend process for `backendId`. Tries three sources + * in order: + * 1. An already-running proc owned by the manager (second-and-later + * sessions on the same backend). + * 2. An in-flight spawn — concurrent callers join the first one. + * 3. The preloader's warm probe (skips spawn + `initialize` handshake). + * 4. A fresh `descriptor.createBackendProcess()` + `start()`. + * + * The returned `warm` slot is non-null only on path 3 and carries the + * probe session id + state snapshot so callers that open a fresh chat + * can adopt the probe session as the user's session (avoids a second + * `newSession` round-trip). Callers that need a specific saved session + * (e.g. `tryResumeSessionFromHistory`) can ignore `warm` — the probe + * session simply sits unused on the proc. + */ private async ensureBackend( backendId: BackendId, descriptor: BackendDescriptor - ): Promise { + ): Promise<{ proc: BackendProcess; warm: WarmBackend | null }> { const existing = this.backends.get(backendId); - if (existing && existing.isRunning()) return existing; + if (existing && existing.isRunning()) return { proc: existing, warm: null }; const inflight = this.starting.get(backendId); - if (inflight) return inflight; + if (inflight) return { proc: await inflight, warm: null }; + + const warm = this.preloader.takeWarm(backendId); + if (warm) { + // Probe subprocess is already started + initialize-handshaken — + // wire it into the manager without paying either cost again. + warm.proc.setPermissionPrompter(this.opts.permissionPrompter); + this.installBackendExitHandler(backendId, warm.proc, descriptor); + this.backends.set(backendId, warm.proc); + return { proc: warm.proc, warm }; + } const proc = descriptor.createBackendProcess({ plugin: this.plugin, @@ -974,43 +1056,56 @@ export class AgentSessionManager { // initialize handshake. In-process adapters (Claude SDK) omit it. if (proc.start) await proc.start(); proc.setPermissionPrompter(this.opts.permissionPrompter); - proc.onExit(() => { - // Backend died unexpectedly. Sessions belonging to *this* backend - // are now unusable (their backend session ids are dead) — but other - // backends keep running. Preserving message history across crashes - // is M5. - if (this.backends.get(backendId) === proc) this.backends.delete(backendId); - const dead = Array.from(this.sessions.values()).filter((s) => s.backendId === backendId); - if (dead.length === 0) return; - for (const s of dead) { - this.detachAutoSave(s.internalId); - this.sessions.delete(s.internalId); - this.chatUIStates.delete(s.internalId); - s.cancel().catch(() => {}); - s.dispose().catch(() => {}); - } - if (this.activeSessionId && !this.sessions.has(this.activeSessionId)) { - const remaining = Array.from(this.sessions.keys()); - this.activeSessionId = remaining[0] ?? null; - } - // Surface the crash so the empty-state pill shows it and the - // router's auto-spawn effect (which bails on lastError) doesn't - // immediately respawn behind the user's back. The next explicit - // create call clears it. - this.lastError = `${descriptor.displayName} backend exited unexpectedly.`; - this.notify(); - }); + this.installBackendExitHandler(backendId, proc, descriptor); this.backends.set(backendId, proc); return proc; })(); this.starting.set(backendId, startPromise); try { - return await startPromise; + return { proc: await startPromise, warm: null }; } finally { this.starting.delete(backendId); } } + /** + * Wire the "backend died unexpectedly" cleanup for `proc`: drop every + * session bound to `backendId`, surface a `lastError`, and notify. Same + * shape for both warm-adopted and freshly-spawned procs. + */ + private installBackendExitHandler( + backendId: BackendId, + proc: BackendProcess, + descriptor: BackendDescriptor + ): void { + proc.onExit(() => { + // Backend died unexpectedly. Sessions belonging to *this* backend + // are now unusable (their backend session ids are dead) — but other + // backends keep running. Preserving message history across crashes + // is M5. + if (this.backends.get(backendId) === proc) this.backends.delete(backendId); + const dead = Array.from(this.sessions.values()).filter((s) => s.backendId === backendId); + if (dead.length === 0) return; + for (const s of dead) { + this.detachAutoSave(s.internalId); + this.sessions.delete(s.internalId); + this.chatUIStates.delete(s.internalId); + s.cancel().catch(() => {}); + s.dispose().catch(() => {}); + } + if (this.activeSessionId && !this.sessions.has(this.activeSessionId)) { + const remaining = Array.from(this.sessions.keys()); + this.activeSessionId = remaining[0] ?? null; + } + // Surface the crash so the empty-state pill shows it and the + // router's auto-spawn effect (which bails on lastError) doesn't + // immediately respawn behind the user's back. The next explicit + // create call clears it. + this.lastError = `${descriptor.displayName} backend exited unexpectedly.`; + this.notify(); + }); + } + /** Whether any session for `backendId` is not safe to dispose yet. */ private hasBusySession(backendId: BackendId): boolean { return Array.from(this.sessions.values()).some((session) => { diff --git a/src/agentMode/ui/AgentModeChat.tsx b/src/agentMode/ui/AgentModeChat.tsx index 16e252d6..2bceba8f 100644 --- a/src/agentMode/ui/AgentModeChat.tsx +++ b/src/agentMode/ui/AgentModeChat.tsx @@ -39,14 +39,18 @@ export const AgentModeChat: React.FC = ({ // Manager fires `notify()` on preload settle, which bumps `tick` above and // re-renders this component — so we can read the flag directly each render. - const preloadReady = manager?.isPreloadReady() ?? true; + // Gate only on the *active* backend so a slow non-active backend never + // holds the chat hostage; the picker shows per-backend loading rows for + // the others. + const preloadReady = manager?.isPreloadReady(descriptor.id) ?? true; // Auto-spawn the first session on mount. The manager de-dupes concurrent // creators via creatingSession, so this is safe to fire whenever the // dependencies change. Skip if the backend isn't installed (the install // pill takes over), there's a prior boot error (Retry handles it), or - // preload hasn't settled (the SDK catalog isn't in cache yet — kicking - // off `newSession` would trigger a redundant on-demand probe). + // the active backend's preload hasn't settled (its catalog isn't in + // cache yet — kicking off `newSession` would trigger a redundant + // on-demand probe). React.useEffect(() => { if (!manager) return; if (!preloadReady) return; diff --git a/src/agentMode/ui/agentModelPickerHelpers.ts b/src/agentMode/ui/agentModelPickerHelpers.ts index ba82e474..57bcaa78 100644 --- a/src/agentMode/ui/agentModelPickerHelpers.ts +++ b/src/agentMode/ui/agentModelPickerHelpers.ts @@ -101,6 +101,27 @@ export function synthesizeAgentEntry( }; } +/** + * Synthesize a non-selectable placeholder row for a backend whose preload + * hasn't settled yet (status `"pending"`) or failed (status `"error"`). + * `_disabledReason` is the contract `ModelSelector` / `ModelEffortPicker` + * use to disable the row and display the right-side label. + */ +export function synthesizePreloadPlaceholder( + descriptor: BackendDescriptor, + status: "pending" | "error" +): ModelSelectorEntry { + const pending = status === "pending"; + return { + ...synthesizeAgentEntry( + `__preload_${status}__`, + pending ? "Loading models…" : "Failed to load", + descriptor + ), + _disabledReason: pending ? "Loading…" : "Unavailable", + }; +} + /** * Resolve a picker entry to its baseModelId. All picker entries from agent * backends are synthesized — the baseModelId lives in `name`. @@ -167,11 +188,23 @@ export function buildPickerEntries( const keepBaseModelId = isActiveBackend ? (ctx.activeModelState?.current.baseModelId ?? null) : (manager.getDefaultSelection(descriptor.id)?.baseModelId ?? null); + const backendModels = cached?.model?.availableModels ?? null; appendBackendSection(entries, descriptor, { - backendModels: cached?.model?.availableModels ?? null, + backendModels, overrides: getBackendModelOverrides(settings, descriptor.id), keepBaseModelId, }); + // No catalog cached yet (and not because the agent intentionally + // omitted a model state). Show a per-backend loading / failure row so + // the user can see every installed backend immediately — important + // because the chat now unblocks on just the *active* backend's + // preload, not the global preload. + if (!backendModels) { + const status = manager.getPreloadStatus(descriptor.id); + if (status === "pending" || status === "error") { + entries.push(synthesizePreloadPlaceholder(descriptor, status)); + } + } } let valueKey = ""; diff --git a/src/agentMode/ui/useAgentModelPicker.ts b/src/agentMode/ui/useAgentModelPicker.ts index c2d37b78..911495cc 100644 --- a/src/agentMode/ui/useAgentModelPicker.ts +++ b/src/agentMode/ui/useAgentModelPicker.ts @@ -67,7 +67,14 @@ function useAgentModelSignal( session?.hasUserVisibleMessages() ? "1" : "0", ]; for (const d of descriptors) { - parts.push(`${d.id}:${modelStateSignature(manager.getCachedBackendState(d.id))}`); + // Include preload status so the picker re-renders when a backend + // flips pending → ready (the placeholder row swaps out for real + // models) without waiting on an unrelated cache write. + parts.push( + `${d.id}:${manager.getPreloadStatus(d.id)}:${modelStateSignature( + manager.getCachedBackendState(d.id) + )}` + ); } return parts.join("|"); }, [manager, descriptors]);