diff --git a/mobile/src/session/mobile-host-model-catalog-query.test.ts b/mobile/src/session/mobile-host-model-catalog-query.test.ts new file mode 100644 index 00000000000..2abb4b8847a --- /dev/null +++ b/mobile/src/session/mobile-host-model-catalog-query.test.ts @@ -0,0 +1,78 @@ +import { createElement, useLayoutEffect } from 'react' +import { act, type ReactTestRenderer } from 'react-test-renderer' +import { describe, expect, it } from 'vitest' +import type { AgentSessionModelCatalogResult } from '../../../src/shared/agent-session/agent-session-wire' +import { AGENT_SESSION_MODEL_CATALOG_SAVED_ONLY_RUNTIME_CAPABILITY } from '../../../src/shared/agent-session/structured-agent-session-surface-capabilities' +import { createWithQueryClient } from '../test-support/query-client-test-renderer' +import type { RpcClient } from '../transport/rpc-client' +import { recordRuntimeStatusProbe } from '../transport/runtime-status-probe' +import { useHostModelCatalogCache } from './mobile-host-model-catalog-query' +import type { HostModelCatalogCache } from './use-mobile-host-model-catalog-upgrade' + +function listed(id: string): AgentSessionModelCatalogResult { + return { + origin: 'probe', + models: [{ id, label: id, isDefault: true, efforts: [] }], + fetchedAt: 1, + listingNamesConfiguredModel: false + } +} + +function heldPreloadClient() { + let answer: (catalog: AgentSessionModelCatalogResult) => void = () => {} + const client: RpcClient = { + sendRequest: () => + new Promise((resolve) => { + answer = (catalog) => + resolve({ id: 'r', ok: true, result: catalog, _meta: { runtimeId: 'runtime-1' } }) + }), + subscribe: () => () => {}, + updateTerminalSubscriptionViewport: () => {}, + getState: () => 'connected', + getReconnectAttempt: () => 0, + getLastConnectedAt: () => 1, + onStateChange: () => () => {}, + notifyForeground: () => {}, + close: () => {} + } + recordRuntimeStatusProbe(client, { + capabilities: [AGENT_SESSION_MODEL_CATALOG_SAVED_ONLY_RUNTIME_CAPABILITY] + }) + return { client, answer: (catalog: AgentSessionModelCatalogResult) => answer(catalog) } +} + +async function mountCache(client: RpcClient) { + const mounted: { cache: HostModelCatalogCache | null } = { cache: null } + function Probe(): null { + const current = useHostModelCatalogCache({ + hostId: 'host-1', + client, + agent: 'claude', + enabled: true + }) + useLayoutEffect(() => { + mounted.cache = current + }) + return null + } + let renderer: ReactTestRenderer | null = null + await act(async () => { + renderer = createWithQueryClient(createElement(Probe)) + }) + return { mounted, unmount: () => act(async () => renderer?.unmount()) } +} + +describe('useHostModelCatalogCache', () => { + it("keeps a new chat's answer that lands while an older preload is out", async () => { + const host = heldPreloadClient() + const { mounted, unmount } = await mountCache(host.client) + + await act(async () => mounted.cache?.record('wt-1', listed('opus'))) + await act(async () => host.answer(listed('sonnet'))) + + expect(mounted.cache?.seed({ newLaunch: true, worktree: 'wt-1', seedsModel: false })).toEqual( + listed('opus') + ) + await unmount() + }) +}) diff --git a/mobile/src/session/mobile-host-model-catalog-query.ts b/mobile/src/session/mobile-host-model-catalog-query.ts index 93b56929476..a4bfe515754 100644 --- a/mobile/src/session/mobile-host-model-catalog-query.ts +++ b/mobile/src/session/mobile-host-model-catalog-query.ts @@ -72,11 +72,22 @@ export function useHostModelCatalogCache(args: { launch ), // A new chat's answer is the account's a new chat pins, so the next chat starts from it. - record: (worktree, catalog) => - queryClient.setQueryData( - queryKey, - (snapshot) => withHostModelCatalogAnswer(snapshot ?? undefined, worktree, catalog) ?? null - ) + record: (worktree, catalog) => { + const held = + queryClient.getQueryData(queryKey) ?? undefined + const next = withHostModelCatalogAnswer(held, worktree, catalog) + if (next === held) { + return + } + queryClient.setQueryData(queryKey, next ?? null) + // A preload sent before this list is older; cancelling it keeps the list just set. + if ( + next?.byWorktree.has(worktree) && + queryClient.isFetching({ queryKey, exact: true }) > 0 + ) { + void queryClient.cancelQueries({ queryKey, exact: true }, { silent: true }) + } + } } }, [agent, exec, hostId, queryClient]) } diff --git a/mobile/src/session/use-mobile-host-model-catalog-upgrade.ts b/mobile/src/session/use-mobile-host-model-catalog-upgrade.ts index 3fd44f74c26..72a3a902735 100644 --- a/mobile/src/session/use-mobile-host-model-catalog-upgrade.ts +++ b/mobile/src/session/use-mobile-host-model-catalog-upgrade.ts @@ -34,6 +34,8 @@ export function useMobileHostModelCatalogUpgrade(args: { newLaunch: boolean /** Where a new chat runs: the host names no default its config could replace. */ worktree?: string + /** The chat has a turn or a send, so its answers are its own account's, which a new chat may no longer pin. */ + hasTurn?: boolean fence: number | null optionCatalog: AgentSessionOptionCatalog | null receiveCatalog: ( @@ -62,9 +64,12 @@ export function useMobileHostModelCatalogUpgrade(args: { worktree } = args const receiver = useRef(args.receiveCatalog) + const recordsNewChat = useRef(newLaunch && args.hasTurn !== true) + const recordingWait = useRef(null) useEffect(() => { receiver.current = args.receiveCatalog - }, [args.receiveCatalog]) + recordsNewChat.current = newLaunch && args.hasTurn !== true + }, [args.receiveCatalog, newLaunch, args.hasTurn]) const [rereads, setRereads] = useState(0) const selectedModel = usesModelSpecificAcpEfforts(agent ?? '') ? args.selectedModel : undefined const recheck = useCallback(() => setRereads((count) => count + 1), []) @@ -73,6 +78,12 @@ export function useMobileHostModelCatalogUpgrade(args: { return } let stale = false + const key = `${agent}\u0000${sessionId}` + // A read sent before the chat's first turn or send answers for the account a new chat pins, + // and so does the listing it waits on, whichever later run picks that up. + const records = + recordsNewChat.current || + (recordingWait.current === key && isHostModelListingWaitInFlight(key, client)) const params = { agent, sessionId, @@ -98,7 +109,7 @@ export function useMobileHostModelCatalogUpgrade(args: { return } // A new chat's answer is the account's a new chat pins, so the next chat starts from it. - if (newLaunch) { + if (records) { modelCatalogCache?.record(worktree ?? '', catalog) } receiver.current(catalog, settled ? 'settled' : event) @@ -109,9 +120,9 @@ export function useMobileHostModelCatalogUpgrade(args: { receiver.current(null, 'demand') } } - const key = `${agent}\u0000${sessionId}` let leave: (() => void) | undefined const wait = () => { + recordingWait.current = records ? key : null leave = joinHostModelListingWait( key, () => read(true), diff --git a/mobile/src/session/use-mobile-structured-agent-options.ts b/mobile/src/session/use-mobile-structured-agent-options.ts index 5a4af847ce8..4e5e5403306 100644 --- a/mobile/src/session/use-mobile-structured-agent-options.ts +++ b/mobile/src/session/use-mobile-structured-agent-options.ts @@ -187,6 +187,7 @@ export function useMobileStructuredAgentOptions(args: { fence, newLaunch: createdHere !== undefined, ...(createdHere ? { worktree: createdHere.worktree } : {}), + hasTurn: args.hasTurn, optionCatalog, activeOptionRecordRef, updateOptionState, diff --git a/mobile/src/session/use-mobile-structured-option-state.ts b/mobile/src/session/use-mobile-structured-option-state.ts index d25cec5f411..db8bac20822 100644 --- a/mobile/src/session/use-mobile-structured-option-state.ts +++ b/mobile/src/session/use-mobile-structured-option-state.ts @@ -11,10 +11,7 @@ import { BoundedMap } from '../../../src/shared/bounded-map' import type { RpcClient } from '../transport/rpc-client' import type { AgentSessionOptionCatalog } from '../../../src/shared/agent-session/options/agent-session-option-catalog-types' import { usesInstalledAgentModelCatalog } from '../../../src/shared/agent-session/options/agent-model-catalog-ownership' -import { - applyStructuredAgentSessionModelCatalog, - createStructuredAgentSessionOptionState -} from '../../../src/shared/agent-session/options/structured-agent-session-options' +import { structuredAgentSessionFirstFrameOptionState } from '../../../src/shared/agent-session/options/structured-agent-session-first-frame' import type { HostModelCatalogCache } from './use-mobile-host-model-catalog-upgrade' import { structuredAgentSessionReportsAtFence } from '../../../src/shared/agent-session/options/structured-agent-session-runtime-reports' @@ -70,18 +67,13 @@ export function useMobileStructuredOptionState(args: { const providerStarting = summary?.hostExecutionPhase === 'starting' const coherentCatalog = usesInstalledAgentModelCatalog(agent ?? '', hostCoherentCatalog) // The host's last list for this agent when this world holds one, else the quiet placeholder until the host answers. - const firstFrame = () => { - const seeded = createStructuredAgentSessionOptionState(agent ?? 'codex', optionCatalog) - const listed = - optionCatalog && - modelCatalogCache?.seed({ newLaunch, seedsModel: false, ...(worktree ? { worktree } : {}) }) - return listed && optionCatalog - ? applyStructuredAgentSessionModelCatalog(seeded, optionCatalog, listed, { - newLaunch, - coherentCatalog - }) - : seeded - } + const firstFrame = () => + structuredAgentSessionFirstFrameOptionState( + agent ?? 'codex', + optionCatalog, + modelCatalogCache?.seed({ newLaunch, seedsModel: false, ...(worktree ? { worktree } : {}) }), + { newLaunch, coherentCatalog } + ) // One view serves every chat tab, so each chat keeps its own reader and its last answer on return. const [readers] = useState( () => diff --git a/src/renderer/src/components/native-chat/session/options/catalog/use-host-model-catalog-upgrade.ts b/src/renderer/src/components/native-chat/session/options/catalog/use-host-model-catalog-upgrade.ts index 6f6c6645d11..15d80316c1d 100644 --- a/src/renderer/src/components/native-chat/session/options/catalog/use-host-model-catalog-upgrade.ts +++ b/src/renderer/src/components/native-chat/session/options/catalog/use-host-model-catalog-upgrade.ts @@ -38,6 +38,8 @@ export function useHostModelCatalogUpgrade(args: { newLaunch: boolean /** Where the launch runs: the host names no default its config could replace. */ worktree?: string + /** The chat has a turn or a send, so its answers are its own account's, which a new chat may no longer pin. */ + hasTurn?: boolean fence: number | null /** The chat's running turn: one running proves its start, which the host re-checks against. */ turnId?: string | null @@ -70,9 +72,12 @@ export function useHostModelCatalogUpgrade(args: { worktree } = args const receiver = useRef(args.receiveCatalog) + const recordsNewChat = useRef(newLaunch && args.hasTurn !== true) + const recordingWait = useRef(null) useEffect(() => { receiver.current = args.receiveCatalog - }, [args.receiveCatalog]) + recordsNewChat.current = newLaunch && args.hasTurn !== true + }, [args.receiveCatalog, newLaunch, args.hasTurn]) const waitKey = `${structuredAgentSessionHostKey(target)}\u0000${agent}\u0000${sessionId}` const [verdict, setVerdict] = useState<{ key: string @@ -104,6 +109,11 @@ export function useHostModelCatalogUpgrade(args: { return } let stale = false + // A read sent before the chat's first turn or send answers for the account a new chat pins, + // and so does the listing it waits on, whichever later run picks that up. + const records = + recordsNewChat.current || + (recordingWait.current === waitKey && isHostModelListingWaitInFlight(waitKey)) const params = { agent, sessionId, @@ -134,13 +144,14 @@ export function useHostModelCatalogUpgrade(args: { : { key: waitKey, unavailable: next } }) // A new chat's answer is the account's a new chat pins, so the next chat starts from it. - if (catalog && newLaunch) { + if (catalog && records) { recordHostModelCatalogSnapshot(target, agent, worktree ?? '', catalog) } receiver.current(catalog, settled ? 'settled' : event) } let leave: (() => void) | null = null const waitForListing = (): void => { + recordingWait.current = records ? waitKey : null leave = joinHostModelListingWait( waitKey, () => read(true), diff --git a/src/renderer/src/components/native-chat/session/options/structured/structured-agent-session-option-first-frame.ts b/src/renderer/src/components/native-chat/session/options/structured/structured-agent-session-option-first-frame.ts index 73864558918..4f3d67bd44a 100644 --- a/src/renderer/src/components/native-chat/session/options/structured/structured-agent-session-option-first-frame.ts +++ b/src/renderer/src/components/native-chat/session/options/structured/structured-agent-session-option-first-frame.ts @@ -1,22 +1,20 @@ import type { AgentType } from '../../../../../../../shared/agent-status/contracts/agent-status-types' import type { RuntimeClientTarget } from '../../../../../../../shared/runtime-client-target' import { structuredAgentSessionSeedCatalog } from '../../../../../../../shared/agent-session/options/structured-agent-session-seed-catalog' -import { - applyStructuredAgentSessionModelCatalog, - createStructuredAgentSessionOptionState, - type StructuredAgentSessionOptionState -} from '../../../../../../../shared/agent-session/options/structured-agent-session-options' +import type { StructuredAgentSessionOptionState } from '../../../../../../../shared/agent-session/options/structured-agent-session-options' +import { structuredAgentSessionFirstFrameOptionState } from '../../../../../../../shared/agent-session/options/structured-agent-session-first-frame' import { readHostModelCatalogSnapshot } from '@/runtime/host-model-catalog-snapshots' -/** The first frame: the host's saved list when the renderer already holds it, else the seed, - * which shows the quiet placeholder until the host answers. */ +/** The first frame from the host's saved list when the renderer already holds it. */ export function firstFrameOptionState( agent: AgentType, target: RuntimeClientTarget, launch: { newLaunch: boolean; worktree?: string; seedsModel: boolean } ): StructuredAgentSessionOptionState { - const seed = structuredAgentSessionSeedCatalog(agent) - const seeded = createStructuredAgentSessionOptionState(agent, seed) - const snapshot = readHostModelCatalogSnapshot(target, agent, launch) - return snapshot ? applyStructuredAgentSessionModelCatalog(seeded, seed, snapshot, launch) : seeded + return structuredAgentSessionFirstFrameOptionState( + agent, + structuredAgentSessionSeedCatalog(agent), + readHostModelCatalogSnapshot(target, agent, launch), + launch + ) } diff --git a/src/renderer/src/components/native-chat/session/options/structured/use-structured-agent-session-options.ts b/src/renderer/src/components/native-chat/session/options/structured/use-structured-agent-session-options.ts index 175bd9a2947..0a192e5b374 100644 --- a/src/renderer/src/components/native-chat/session/options/structured/use-structured-agent-session-options.ts +++ b/src/renderer/src/components/native-chat/session/options/structured/use-structured-agent-session-options.ts @@ -95,6 +95,7 @@ export function useStructuredAgentSessionOptions(args: StructuredAgentSessionOpt // A resumed conversation may keep its own model, so only a new one runs the listed default. newLaunch, ...(launch?.worktree ? { worktree: launch.worktree } : {}), + hasTurn: args.hasTurn, fence, turnId, reportedUnpickedModel: diff --git a/src/renderer/src/runtime/host-model-catalog-snapshots.test.ts b/src/renderer/src/runtime/host-model-catalog-snapshots.test.ts index 0e5a303cc5d..07a7ee10923 100644 --- a/src/renderer/src/runtime/host-model-catalog-snapshots.test.ts +++ b/src/renderer/src/runtime/host-model-catalog-snapshots.test.ts @@ -31,6 +31,7 @@ import { getDefaultSettings } from '../../../shared/constants' import type { PublicKnownRuntimeEnvironment } from '../../../shared/remote-runtime/host-catalog/runtime-environments' import { useAppStore } from '@/store' import { + preloadHostModelCatalogSnapshots, readHostModelCatalogSnapshot, recordHostModelCatalogSnapshot, resetHostModelCatalogSnapshotsForTests @@ -83,6 +84,31 @@ describe('host model catalog snapshots', () => { recordHostModelCatalogSnapshot(LOCAL, 'grok', 'wt-1', { origin: 'unknown' }) expect(readHostModelCatalogSnapshot(LOCAL, 'grok', NEW_CHAT)).toBe(worldwide) }) + + it('keeps a held list through an answer whose listing is still running', () => { + const held = list(true) + recordHostModelCatalogSnapshot(LOCAL, 'claude', 'wt-1', held) + recordHostModelCatalogSnapshot(LOCAL, 'claude', 'wt-1', { + origin: 'unknown', + listingInProgress: true + }) + expect(readHostModelCatalogSnapshot(LOCAL, 'claude', NEW_CHAT)).toBe(held) + const other = { ...NEW_CHAT, worktree: 'wt-2' } + const filling = { ...list(true), listingInProgress: true as const } + recordHostModelCatalogSnapshot(LOCAL, 'claude', 'wt-2', filling) + expect(readHostModelCatalogSnapshot(LOCAL, 'claude', other)).toBe(filling) + }) + + it("keeps a chat's answer that lands while an older preload is out", async () => { + let answer: (catalog: AgentSessionModelCatalogResult) => void = () => {} + mocks.call.mockImplementation(() => new Promise((resolve) => (answer = resolve))) + const preload = preloadHostModelCatalogSnapshots(LOCAL, ['claude'], { savedOnly: true }) + const chat = list(true) + recordHostModelCatalogSnapshot(LOCAL, 'claude', 'wt-1', chat) + answer(list(false)) + await preload + expect(readHostModelCatalogSnapshot(LOCAL, 'claude', NEW_CHAT)).toBe(chat) + }) }) const SAVED_ONLY = 'agent-session.model-catalog.saved-only.v1' diff --git a/src/renderer/src/runtime/host-model-catalog-snapshots.ts b/src/renderer/src/runtime/host-model-catalog-snapshots.ts index c5fc9357e58..a53ea6306dc 100644 --- a/src/renderer/src/runtime/host-model-catalog-snapshots.ts +++ b/src/renderer/src/runtime/host-model-catalog-snapshots.ts @@ -13,6 +13,8 @@ import { structuredAgentSessionHostKey } from './structured-agent-session-host-c const snapshots = new Map() // Bumped when a host's answers are dropped, so a preload sent before then records nothing. const generations = new Map() +// Bumped per host and agent by each chat's answer, so an older preload never replaces it. +const chatAnswers = new Map() function snapshotKey(hostKey: string, agent: string): string { return `${hostKey}\u0000${agent}` @@ -39,6 +41,9 @@ export function recordHostModelCatalogSnapshot( ): void { const key = snapshotKey(structuredAgentSessionHostKey(target), agent) const next = withHostModelCatalogAnswer(snapshots.get(key), worktree, catalog) + if (worktree !== null && next !== snapshots.get(key) && next?.byWorktree.has(worktree)) { + chatAnswers.set(key, (chatAnswers.get(key) ?? 0) + 1) + } if (next) { snapshots.set(key, next) } else { @@ -78,23 +83,29 @@ export async function preloadHostModelCatalogSnapshots( const hostKey = structuredAgentSessionHostKey(target) const generation = generations.get(hostKey) ?? 0 await Promise.all( - agents.map((agent) => - callStructuredAgentSession( + agents.map((agent) => { + const key = snapshotKey(hostKey, agent) + const answersAtSend = chatAnswers.get(key) ?? 0 + return callStructuredAgentSession( target, 'agentSession.modelCatalog', { agent, ...options } ) .then((catalog) => { - if ((generations.get(hostKey) ?? 0) === generation) { + if ( + (generations.get(hostKey) ?? 0) === generation && + (chatAnswers.get(key) ?? 0) === answersAtSend + ) { recordHostModelCatalogSnapshot(target, agent, null, catalog) } }) .catch(() => undefined) - ) + }) ) } export function resetHostModelCatalogSnapshotsForTests(): void { snapshots.clear() generations.clear() + chatAnswers.clear() } diff --git a/src/shared/agent-session/options/host-model-catalog-snapshot-rules.ts b/src/shared/agent-session/options/host-model-catalog-snapshot-rules.ts index 6298b02a7c8..00557f102bb 100644 --- a/src/shared/agent-session/options/host-model-catalog-snapshot-rules.ts +++ b/src/shared/agent-session/options/host-model-catalog-snapshot-rules.ts @@ -53,16 +53,30 @@ export function selectHostModelCatalogSnapshot( } /** The snapshot after a host answer for `worktree` (null for a session-less read); undefined means keep none. - * An unknown or empty answer proves absence and retires what it answered for. */ + * An unknown or empty answer proves absence and retires what it answered for. One whose listing is + * still running proves nothing yet: it never replaces or retires a held list, but may fill an empty scope. */ export function withHostModelCatalogAnswer( snapshot: HostModelCatalogSnapshot | undefined, worktree: string | null, catalog: AgentSessionModelCatalogResult ): HostModelCatalogSnapshot | undefined { - if (catalog.origin === 'unknown' || catalog.models.length === 0) { + const listed = catalog.origin !== 'unknown' && catalog.models.length > 0 + if (catalog.listingInProgress) { + const held = + worktree === null + ? !!snapshot && (snapshot.account !== null || snapshot.byWorktree.size > 0) + : snapshot?.byWorktree.has(worktree) === true + if (held || !listed) { + return snapshot + } + } + if (!listed) { if (worktree === null || !snapshot) { return undefined } + if (!snapshot.byWorktree.has(worktree)) { + return snapshot + } const retired = snapshot.byWorktree.get(worktree) const byWorktree = new Map(snapshot.byWorktree) byWorktree.delete(worktree) diff --git a/src/shared/agent-session/options/structured-agent-session-first-frame.ts b/src/shared/agent-session/options/structured-agent-session-first-frame.ts new file mode 100644 index 00000000000..be9534cf7ce --- /dev/null +++ b/src/shared/agent-session/options/structured-agent-session-first-frame.ts @@ -0,0 +1,21 @@ +import type { AgentSessionModelCatalogResult } from '../agent-session-wire' +import type { AgentSessionOptionCatalog } from './agent-session-option-catalog-types' +import { + applyStructuredAgentSessionModelCatalog, + createStructuredAgentSessionOptionState, + type StructuredAgentSessionOptionState +} from './structured-agent-session-options' + +/** A chat's first frame: the host's held list when the client has one, else the seed, which shows + * the quiet placeholder until the host answers. */ +export function structuredAgentSessionFirstFrameOptionState( + agent: string, + seed: AgentSessionOptionCatalog | null, + held: AgentSessionModelCatalogResult | undefined, + options: { newLaunch: boolean; coherentCatalog?: boolean } +): StructuredAgentSessionOptionState { + const seeded = createStructuredAgentSessionOptionState(agent, seed) + return held && seed + ? applyStructuredAgentSessionModelCatalog(seeded, seed, held, options) + : seeded +}