Fix provider snapshot session hydration and agent scoping

This commit is contained in:
Mohamed Boudra
2026-04-07 00:11:43 +07:00
parent 27ddc95862
commit 390a3402ab
5 changed files with 49 additions and 211 deletions
@@ -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<string, AgentModelDefinition[]>();
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 ||
@@ -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;
+14 -187
View File
@@ -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<string | undefined>(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<AgentProvider, AgentProviderDefinition>(
@@ -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<AgentProvider, AgentProviderDefinition>(
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<string, AgentModelDefinition[]>();
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();
@@ -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,
};
@@ -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) =>