Fix provider refresh scope and self-healing

packages/app/src/hooks/use-providers-snapshot.ts

@@ -6,7 +6,7 @@ import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-
 import { useSessionStore } from "@/stores/session-store";
 import { queryClient as singletonQueryClient } from "@/query/query-client";
 
-function normalizeProvidersSnapshotCwdKey(cwd?: string | null): string | null {
+export function normalizeProvidersSnapshotCwdKey(cwd?: string | null): string | null {
   const trimmed = cwd?.trim();
   if (!trimmed) {
     return null;
@@ -15,10 +15,33 @@ function normalizeProvidersSnapshotCwdKey(cwd?: string | null): string | null {
   return trimmed.replace(/^\/(?:Users|home)\/[^/]+/, "~");
 }
 
+export function shouldApplyProvidersSnapshotUpdate(
+  currentCwd?: string | null,
+  messageCwd?: string | null,
+): boolean {
+  const currentCwdKey = normalizeProvidersSnapshotCwdKey(currentCwd);
+  const messageCwdKey = normalizeProvidersSnapshotCwdKey(messageCwd);
+  return messageCwdKey === currentCwdKey || (currentCwdKey === null && messageCwdKey === "~");
+}
+
 export function providersSnapshotQueryKey(serverId: string | null, cwd?: string | null) {
   return ["providersSnapshot", serverId, normalizeProvidersSnapshotCwdKey(cwd)] as const;
 }
 
+function providersSnapshotRequest(cwd?: string): { cwd?: string } {
+  return cwd ? { cwd } : {};
+}
+
+function refreshProvidersSnapshotRequest(
+  cwd: string | undefined,
+  providers: AgentProvider[] | undefined,
+): { cwd?: string; providers?: AgentProvider[] } {
+  return {
+    ...providersSnapshotRequest(cwd),
+    ...(providers ? { providers } : {}),
+  };
+}
+
 interface UseProvidersSnapshotResult {
   entries: ProviderSnapshotEntry[] | undefined;
   isLoading: boolean;
@@ -62,7 +85,7 @@ export function useProvidersSnapshot(
       if (!client) {
         throw new Error("Host is not connected");
       }
-      return client.getProvidersSnapshot({ cwd: normalizedCwd });
+      return client.getProvidersSnapshot(providersSnapshotRequest(normalizedCwd));
     },
   });
 
@@ -71,7 +94,9 @@ export function useProvidersSnapshot(
       if (!client) {
         return;
       }
-      await client.refreshProvidersSnapshot({ cwd: normalizedCwd, providers });
+      await client.refreshProvidersSnapshot(
+        refreshProvidersSnapshotRequest(normalizedCwd, providers),
+      );
     },
   });
   const { mutateAsync: refreshSnapshot, isPending: isRefreshing } = refreshMutation;
@@ -85,8 +110,7 @@ export function useProvidersSnapshot(
       if (message.type !== "providers_snapshot_update") {
         return;
       }
-      const messageCwdKey = normalizeProvidersSnapshotCwdKey(message.payload.cwd);
-      if (messageCwdKey !== normalizedCwdKey) {
+      if (!shouldApplyProvidersSnapshotUpdate(normalizedCwd, message.payload.cwd)) {
         return;
       }
       queryClient.setQueryData(queryKey, {
@@ -99,6 +123,7 @@ export function useProvidersSnapshot(
     client,
     enabled,
     isConnected,
+    normalizedCwd,
     normalizedCwdKey,
     queryClient,
     queryKey,
@@ -108,9 +133,14 @@ export function useProvidersSnapshot(
 
   const refresh = useCallback(
     async (providers?: AgentProvider[]) => {
+      if (!client) {
+        return;
+      }
       await refreshSnapshot(providers);
+      const snapshot = await client.getProvidersSnapshot(providersSnapshotRequest(normalizedCwd));
+      queryClient.setQueryData(queryKey, snapshot);
     },
-    [refreshSnapshot],
+    [client, normalizedCwd, queryClient, queryKey, refreshSnapshot],
   );
 
   const refetchIfStale = useCallback(() => {
@@ -139,6 +169,6 @@ export function prefetchProvidersSnapshot(
   void singletonQueryClient.prefetchQuery({
     queryKey,
     staleTime: 60_000,
-    queryFn: () => client.getProvidersSnapshot({ cwd: normalizedCwd }),
+    queryFn: () => client.getProvidersSnapshot(providersSnapshotRequest(normalizedCwd)),
   });
 }