From 390a3402abdb27f8cf9cf3edd48df376bc713b5c Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 7 Apr 2026 00:11:43 +0700 Subject: [PATCH] Fix provider snapshot session hydration and agent scoping --- .../app/src/components/agent-status-bar.tsx | 17 +- packages/app/src/contexts/session-context.tsx | 29 +++ .../app/src/hooks/use-agent-form-state.ts | 201 ++---------------- .../app/src/hooks/use-providers-snapshot.ts | 2 + .../screens/workspace/workspace-screen.tsx | 11 - 5 files changed, 49 insertions(+), 211 deletions(-) diff --git a/packages/app/src/components/agent-status-bar.tsx b/packages/app/src/components/agent-status-bar.tsx index bd273c15a..b1cf99aa4 100644 --- a/packages/app/src/components/agent-status-bar.tsx +++ b/packages/app/src/components/agent-status-bar.tsx @@ -887,26 +887,17 @@ export function AgentStatusBar({ agentId, serverId, onDropdownClose }: AgentStat const models = snapshotModels; const agentProviderDefinitions = useMemo(() => { - if (snapshotEntries) { - return AGENT_PROVIDER_DEFINITIONS.filter((d) => - snapshotEntries.some((e) => e.provider === d.id), - ); - } const definition = AGENT_PROVIDER_DEFINITIONS.find((d) => d.id === agent?.provider); return definition ? [definition] : []; - }, [agent?.provider, snapshotEntries]); + }, [agent?.provider]); const agentProviderModels = useMemo(() => { const map = new Map(); - if (snapshotEntries) { - for (const entry of snapshotEntries) { - if (entry.models) { - map.set(entry.provider, entry.models); - } - } + if (agent?.provider && snapshotModels) { + map.set(agent.provider, snapshotModels); } return map; - }, [snapshotEntries]); + }, [agent?.provider, snapshotModels]); const displayMode = availableModes.find((mode) => mode.id === agent?.currentModeId)?.label || diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 9c36e8132..a7a2717ac 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -5,6 +5,7 @@ import { useQueryClient } from "@tanstack/react-query"; import { useClientActivity } from "@/hooks/use-client-activity"; import { usePushTokenRegistration } from "@/hooks/use-push-token-registration"; import { clearArchiveAgentPending } from "@/hooks/use-archive-agent"; +import { prefetchProvidersSnapshot } from "@/hooks/use-providers-snapshot"; import { generateMessageId, type StreamItem } from "@/types/stream"; import { processTimelineResponse, @@ -603,6 +604,34 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider updateSessionClient(serverId, client); }, [serverId, client, updateSessionClient]); + useEffect(() => { + const serverInfo = client.getLastServerInfoMessage(); + if (!serverInfo) { + return; + } + + updateSessionServerInfo(serverId, { + serverId: serverInfo.serverId, + hostname: serverInfo.hostname, + version: serverInfo.version, + ...(serverInfo.capabilities ? { capabilities: serverInfo.capabilities } : {}), + ...(serverInfo.features ? { features: serverInfo.features } : {}), + }); + }, [client, serverId, updateSessionServerInfo]); + + useEffect(() => { + if (!isConnected) { + return; + } + + const serverInfo = client.getLastServerInfoMessage(); + if (!serverInfo?.features?.providersSnapshot) { + return; + } + + prefetchProvidersSnapshot(serverId, client); + }, [client, isConnected, serverId]); + useEffect(() => { if (!voiceRuntime) { return; diff --git a/packages/app/src/hooks/use-agent-form-state.ts b/packages/app/src/hooks/use-agent-form-state.ts index 5d5865714..09e0ae2fd 100644 --- a/packages/app/src/hooks/use-agent-form-state.ts +++ b/packages/app/src/hooks/use-agent-form-state.ts @@ -1,5 +1,4 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { useQuery, useQueries } from "@tanstack/react-query"; import { AGENT_PROVIDER_DEFINITIONS, type AgentProviderDefinition, @@ -11,7 +10,6 @@ import type { ProviderSnapshotEntry, } from "@server/server/agent/agent-sdk-types"; import { useHosts } from "@/runtime/host-runtime"; -import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime"; import { useProvidersSnapshot } from "./use-providers-snapshot"; import { useFormPreferences, @@ -371,35 +369,19 @@ export function useAgentFormState(options: UseAgentFormStateOptions = {}): UseAg } }, [isVisible]); - // Session state for provider model listing - const client = useHostRuntimeClient(formState.serverId ?? ""); - const isConnected = useHostRuntimeIsConnected(formState.serverId ?? ""); const { entries: snapshotEntries, isLoading: snapshotIsLoading, isFetching: snapshotIsFetching, - supportsSnapshot: supportsProvidersSnapshot, + error: snapshotError, refresh: refreshSnapshot, } = useProvidersSnapshot(formState.serverId); - const [debouncedCwd, setDebouncedCwd] = useState(undefined); - useEffect(() => { - const trimmed = formState.workingDir.trim(); - const next = trimmed.length > 0 ? trimmed : undefined; - const timer = setTimeout(() => setDebouncedCwd(next), 180); - return () => clearTimeout(timer); - }, [formState.workingDir]); - const allProviderEntries = useMemo( - () => (supportsProvidersSnapshot ? snapshotEntries ?? [] : undefined), - [snapshotEntries, supportsProvidersSnapshot], - ); + const allProviderEntries = useMemo(() => snapshotEntries ?? [], [snapshotEntries]); const snapshotProviderDefinitions = useMemo(() => { - if (!supportsProvidersSnapshot) { - return []; - } const snapshotProviders = new Set((snapshotEntries ?? []).map((entry) => entry.provider)); return allProviderDefinitions.filter((definition) => snapshotProviders.has(definition.id)); - }, [snapshotEntries, supportsProvidersSnapshot]); + }, [snapshotEntries]); const snapshotProviderDefinitionMap = useMemo( () => new Map( @@ -435,157 +417,13 @@ export function useAgentFormState(options: UseAgentFormStateOptions = {}): UseAg snapshotSelectedEntry?.modes ?? snapshotProviderDefinitionMap.get(formState.provider)?.modes ?? []; - - // COMPAT(providersSnapshot): legacy fallback for daemons without snapshot support — remove when all daemons support snapshots - const availableProvidersQuery = useQuery({ - queryKey: ["availableProviders", formState.serverId], - enabled: Boolean( - !supportsProvidersSnapshot && - isVisible && - isTargetDaemonReady && - formState.serverId && - client && - isConnected, - ), - staleTime: 60 * 1000, - queryFn: async () => { - if (!client) { - throw new Error("Host is not connected"); - } - const payload = await client.listAvailableProviders(); - if (payload.error) { - throw new Error(payload.error); - } - return payload.providers.filter((entry) => entry.available).map((entry) => entry.provider); - }, - }); - const legacyProviderDefinitions = useMemo(() => { - const availableProviders = availableProvidersQuery.data; - if (!availableProviders) { - return []; - } - const available = new Set(availableProviders); - return allProviderDefinitions.filter((definition) => available.has(definition.id)); - }, [availableProvidersQuery.data]); - const legacyProviderDefinitionMap = useMemo( - () => - new Map( - legacyProviderDefinitions.map((definition) => [definition.id, definition]), - ), - [legacyProviderDefinitions], - ); - - // COMPAT(providersSnapshot): legacy fallback for daemons without snapshot support — remove when all daemons support snapshots - const legacySelectedProviderModelsQuery = useQuery({ - queryKey: ["providerModels", formState.serverId, formState.provider], - enabled: Boolean( - !supportsProvidersSnapshot && - isVisible && - isTargetDaemonReady && - formState.serverId && - client && - isConnected && - legacyProviderDefinitionMap.has(formState.provider), - ), - staleTime: 5 * 60 * 1000, - queryFn: async () => { - if (!client) { - throw new Error("Host is not connected"); - } - const payload = await client.listProviderModels(formState.provider, { - cwd: debouncedCwd, - }); - if (payload.error) { - throw new Error(payload.error); - } - return payload.models ?? []; - }, - }); - const legacySelectedProviderModels = legacySelectedProviderModelsQuery.data ?? null; - - // COMPAT(providersSnapshot): legacy fallback for daemons without snapshot support — remove when all daemons support snapshots - const providerModesQuery = useQuery({ - queryKey: ["providerModes", formState.serverId, formState.provider, debouncedCwd], - enabled: Boolean( - !supportsProvidersSnapshot && - isVisible && - isTargetDaemonReady && - formState.serverId && - client && - isConnected && - legacyProviderDefinitionMap.has(formState.provider), - ), - staleTime: 5 * 60 * 1000, - queryFn: async () => { - if (!client) { - throw new Error("Host is not connected"); - } - const payload = await client.listProviderModes(formState.provider, { - cwd: debouncedCwd, - }); - if (payload.error) { - throw new Error(payload.error); - } - return payload.modes ?? []; - }, - }); - - // COMPAT(providersSnapshot): legacy fallback for daemons without snapshot support — remove when all daemons support snapshots - const allProviderModelQueries = useQueries({ - queries: (supportsProvidersSnapshot ? [] : legacyProviderDefinitions).map((def) => ({ - queryKey: ["providerModels", formState.serverId, def.id], - enabled: Boolean( - isVisible && isTargetDaemonReady && formState.serverId && client && isConnected, - ), - staleTime: 5 * 60 * 1000, - queryFn: async () => { - if (!client) { - throw new Error("Host is not connected"); - } - const payload = await client.listProviderModels(def.id as AgentProvider, { - cwd: debouncedCwd, - }); - if (payload.error) { - throw new Error(payload.error); - } - return payload.models ?? []; - }, - })), - }); - const legacyAllProviderModels = useMemo(() => { - const map = new Map(); - for (let i = 0; i < legacyProviderDefinitions.length; i++) { - const query = allProviderModelQueries[i]; - if (query?.data) { - map.set(legacyProviderDefinitions[i]!.id, query.data); - } - } - return map; - }, [allProviderModelQueries, legacyProviderDefinitions]); - const legacySelectedProviderModes = - providerModesQuery.data ?? legacyProviderDefinitionMap.get(formState.provider)?.modes ?? []; - - const providerDefinitions = supportsProvidersSnapshot - ? snapshotProviderDefinitions - : legacyProviderDefinitions; - const providerDefinitionMap = supportsProvidersSnapshot - ? snapshotProviderDefinitionMap - : legacyProviderDefinitionMap; - const selectableProviderDefinitionMap = supportsProvidersSnapshot - ? snapshotSelectableProviderDefinitionMap - : legacyProviderDefinitionMap; - const allProviderModels = supportsProvidersSnapshot - ? snapshotAllProviderModels - : legacyAllProviderModels; - const availableModels = supportsProvidersSnapshot - ? snapshotSelectedProviderModels - : legacySelectedProviderModels; - const modeOptions = supportsProvidersSnapshot - ? snapshotSelectedProviderModes - : legacySelectedProviderModes; - const isAllModelsLoading = supportsProvidersSnapshot - ? snapshotIsLoading || snapshotIsFetching - : allProviderModelQueries.some((q) => q.isLoading); + const providerDefinitions = snapshotProviderDefinitions; + const providerDefinitionMap = snapshotProviderDefinitionMap; + const selectableProviderDefinitionMap = snapshotSelectableProviderDefinitionMap; + const allProviderModels = snapshotAllProviderModels; + const availableModels = snapshotSelectedProviderModels; + const modeOptions = snapshotSelectedProviderModes; + const isAllModelsLoading = snapshotIsLoading || snapshotIsFetching; // Combine initialValues with initialServerId for resolution const combinedInitialValues = useMemo((): FormInitialValues | undefined => { @@ -807,12 +645,8 @@ export function useAgentFormState(options: UseAgentFormStateOptions = {}): UseAg }, []); const refreshProviderModels = useCallback(() => { - if (supportsProvidersSnapshot) { - refreshSnapshot(); - return; - } - void legacySelectedProviderModelsQuery.refetch(); - }, [legacySelectedProviderModelsQuery, refreshSnapshot, supportsProvidersSnapshot]); + refreshSnapshot(); + }, [refreshSnapshot]); const persistFormPreferences = useCallback(async () => { const resolvedModel = resolveEffectiveModel(availableModels, formState.model); @@ -848,15 +682,8 @@ export function useAgentFormState(options: UseAgentFormStateOptions = {}): UseAg const effectiveModel = resolveEffectiveModel(availableModels, formState.model); const resolvedModelId = effectiveModel?.id ?? formState.model; const availableThinkingOptions = effectiveModel?.thinkingOptions ?? []; - const isModelLoading = supportsProvidersSnapshot - ? snapshotIsLoading || snapshotIsFetching - : legacySelectedProviderModelsQuery.isLoading || legacySelectedProviderModelsQuery.isFetching; - const modelError = - supportsProvidersSnapshot - ? null - : legacySelectedProviderModelsQuery.error instanceof Error - ? legacySelectedProviderModelsQuery.error.message - : null; + const isModelLoading = snapshotIsLoading || snapshotIsFetching; + const modelError = snapshotError; const workingDirIsEmpty = !formState.workingDir.trim(); diff --git a/packages/app/src/hooks/use-providers-snapshot.ts b/packages/app/src/hooks/use-providers-snapshot.ts index 943791f8a..ffe23fa78 100644 --- a/packages/app/src/hooks/use-providers-snapshot.ts +++ b/packages/app/src/hooks/use-providers-snapshot.ts @@ -14,6 +14,7 @@ interface UseProvidersSnapshotResult { entries: ProviderSnapshotEntry[] | undefined; isLoading: boolean; isFetching: boolean; + error: string | null; supportsSnapshot: boolean; refresh: () => void; } @@ -69,6 +70,7 @@ export function useProvidersSnapshot(serverId: string | null): UseProvidersSnaps entries: snapshotQuery.data?.entries ?? undefined, isLoading: snapshotQuery.isLoading, isFetching: snapshotQuery.isFetching, + error: snapshotQuery.error instanceof Error ? snapshotQuery.error.message : null, supportsSnapshot, refresh, }; diff --git a/packages/app/src/screens/workspace/workspace-screen.tsx b/packages/app/src/screens/workspace/workspace-screen.tsx index dbeb7f930..d4555c4ab 100644 --- a/packages/app/src/screens/workspace/workspace-screen.tsx +++ b/packages/app/src/screens/workspace/workspace-screen.tsx @@ -70,7 +70,6 @@ import { workspaceTabTargetsEqual, } from "@/utils/workspace-tab-identity"; import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime"; -import { prefetchProvidersSnapshot } from "@/hooks/use-providers-snapshot"; import { useWorkspaceTerminalSessionRetention } from "@/terminal/hooks/use-workspace-terminal-session-retention"; import { checkoutStatusQueryKey, @@ -609,16 +608,6 @@ function WorkspaceScreenContent({ serverId, workspaceId }: WorkspaceScreenProps) const client = useHostRuntimeClient(normalizedServerId); const isConnected = useHostRuntimeIsConnected(normalizedServerId); - useEffect(() => { - if (!client || !isConnected || !normalizedServerId) { - return; - } - const session = useSessionStore.getState().sessions[normalizedServerId]; - if (session?.serverInfo?.features?.providersSnapshot) { - prefetchProvidersSnapshot(normalizedServerId, client); - } - }, [client, isConnected, normalizedServerId]); - const workspaceAgentVisibility = useStoreWithEqualityFn( useSessionStore, (state) =>