query-provider.tsx14.1 KBView on GitHub import { hashKey, QueryCache, QueryClient, QueryClientProvider } from '@tanstack/react-query';
import { TRPCProvider, useTRPC, useTRPCClient } from '@/modules/trpc/context';
import { useEffect, useMemo, useRef, useState, type PropsWithChildren } from 'react';
import { createTRPCClient, httpBatchLink } from '@trpc/client';
import { signOut } from '@/modules/auth/utils/auth-client';
import { useAdminViewingUserId } from '@/modules/store';
import type { AppRouter } from '@zero/server/trpc';
import {
clearPersistedQueries,
restoreQueries,
setupQueryPersistence,
} from '@/lib/query-persistence';
import {
installArchivedThreadGuard,
setSuppressionReporter,
} from '@/modules/threads/rendering/archived-thread-guard';
import { getRuntimeBackendUrl } from '@/lib/runtime-urls';
import { trpcRequestIncludes } from '@/lib/trpc-batch-url';
import { trackTrpcRequest } from '@/lib/client-telemetry-trpc';
import { whenSessionSettled } from '@/lib/session-ready';
import { setBrowserQueryClient } from '@/lib/browser-query-client';
import { createDeepDebuggerLink } from '@/modules/debugger/deepDebugger/trpcLink';
import { attachTraceIdToRecent } from '@/modules/debugger/deepDebugger/buffer';
import superjson from 'superjson';
// The join key between a browser action and its server trace, so it must be unique.
// A timestamp plus a counter that resets on page load is not: two tabs, or a reload inside
// the same millisecond, mint identical ids and the join then matches unrelated requests.
function _qpNextBatchId(): string {
// `crypto.randomUUID` is only defined in a secure context (and only from Safari 15.4). This
// runs on the path of EVERY tRPC request, so a throw here would fail the whole app — login
// included — rather than the one feature that elsewhere depends on it. A diagnostic join key
// is never worth that, hence the fallback.
try {
if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') {
return `batch_${crypto.randomUUID()}`;
}
} catch {
// fall through
}
return `batch_${Date.now().toString(36)}_${Math.random().toString(36).slice(2, 10)}`;
}
// Re-export from the stable context module for consumers
export { TRPCProvider, useTRPC, useTRPCClient };
export const makeQueryClient = (connectionId: string | null) =>
new QueryClient({
queryCache: new QueryCache({
onError: (err, { meta }) => {
if (meta && meta.noGlobalError === true) return;
if (meta && typeof meta.customError === 'string') console.error(meta.customError);
else if (
err.message === 'Required scopes missing' ||
err.message.includes('Invalid connection')
) {
signOut({
fetchOptions: {
onSuccess: () => {
if (window.location.href.includes('/login')) return;
window.location.href = '/login?error=required_scopes_missing';
},
},
});
} else console.error(err.message || 'Something went wrong');
},
}),
defaultOptions: {
queries: {
structuralSharing: true, // Prevent re-renders on identical data
staleTime: 1 * 60 * 1000, // 1 minute - balance between freshness and speed
gcTime: 24 * 60 * 60 * 1000,
refetchOnWindowFocus: true,
refetchOnMount: true,
refetchInterval: false,
queryKeyHashFn: (queryKey) => hashKey([{ connectionId }, ...queryKey]),
},
mutations: {
onError: (err) => console.error(err.message),
},
},
});
const browserQueryClient = {
queryClient: null,
activeConnectionId: null,
} as {
queryClient: QueryClient | null;
activeConnectionId: string | null;
};
const getQueryClient = (connectionId: string | null) => {
if (typeof window === 'undefined') {
return makeQueryClient(connectionId);
} else {
if (!browserQueryClient.queryClient || browserQueryClient.activeConnectionId !== connectionId) {
browserQueryClient.queryClient = makeQueryClient(connectionId);
browserQueryClient.activeConnectionId = connectionId;
// Publish it for non-React callers (the store's SSE processors) — see browser-query-client.
setBrowserQueryClient(browserQueryClient.queryClient);
}
return browserQueryClient.queryClient;
}
};
const getUrl = () => getRuntimeBackendUrl() + '/api/trpc';
// Set by forceSendPending() when the page is unloading. Teardown and keepalive are SEPARATE
// concerns: every teardown send must skip the session gate below, because awaiting anything
// while the page is going away is precisely how the request gets lost — but only a send small
// enough for the ~64KB keepalive cap can actually use keepalive, which rules out attachments.
// Conflating them meant an attachment send on teardown still awaited the gate.
let _teardownNextRequest: { keepalive: boolean } | null = null;
export function setTeardownNextRequest(opts: { keepalive: boolean }) {
_teardownNextRequest = opts;
}
// httpBatchLink puts a comma-joined procedure list in the path, so this has to parse the batch
// rather than match the whole pathname — see lib/trpc-batch-url.ts.
function isMailSendTrpcRequest(url: string): boolean {
return trpcRequestIncludes(url, 'mail.send');
}
// Shared fetch handler with redirect logic
const createFetchHandler = () => async (url: string, options?: RequestInit) => {
const batchId = _qpNextBatchId();
const headers = new Headers(options?.headers);
headers.set('X-Client-Request-Id', batchId);
const teardown = _teardownNextRequest && isMailSendTrpcRequest(url) ? _teardownNextRequest : null;
if (teardown) {
_teardownNextRequest = null;
}
const shouldUseKeepalive = teardown?.keepalive === true;
// Hold the request until better-auth has settled whether there's a session.
// Nothing upstream does this: SessionGuard is unmounted and AuthMonitor is a
// non-blocking sibling, so the app shell mounts and fires its whole boot
// fan-out before the session cookie exists — a burst of requests that all
// 401, except public `settings.get`, which answers 200 with all-off defaults
// and tells entitled users their account has no AI. See lib/session-ready.
//
// The transport is the only chokepoint that catches every caller (providers,
// layout hooks, child components alike). Gating renders would not: hooks fire
// their queries regardless of what the component returns.
//
// Skipped for keepalive sends — those happen during page unload, where the
// session settled long ago and awaiting anything risks losing the request.
if (!teardown) {
await whenSessionSettled();
}
// Every tRPC request is recorded against `batchId` — the same id the server sees as
// X-Client-Request-Id — so a server trace and the browser's account of the same call can be
// joined. Aborted and network-failed requests are the reason this exists: they leave no server
// row at all, which is how a send that never left the browser stayed invisible for a day.
// See lib/client-telemetry-trpc.ts.
return trackTrpcRequest({ url, correlationId: batchId, signal: options?.signal }, () =>
fetch(
url,
{
...options,
headers,
credentials: 'include',
...(shouldUseKeepalive ? { keepalive: true } : {}),
},
),
).then(
(res) => {
const currentPath = new URL(window.location.href).pathname;
const redirectPath = res.headers.get('X-Zero-Redirect');
if (!!redirectPath && new URL(redirectPath, window.location.href).pathname !== currentPath) {
window.location.href = redirectPath;
res.headers.delete('X-Zero-Redirect');
}
const traceId = res.headers.get('X-Trace-ID') ?? undefined;
const serverRequestId = res.headers.get('X-Request-ID') ?? undefined;
if (traceId || serverRequestId) {
attachTraceIdToRecent({ traceId, serverRequestId });
}
return res;
},
);
};
// Create tRPC client factory with optional admin header
// Excludes playground routes from admin viewing user mechanism
const createTrpcClient = (adminViewingUserId: string | null) =>
createTRPCClient<AppRouter>({
links: [
// loggerLink({ enabled: () => true }),
createDeepDebuggerLink<AppRouter>(),
httpBatchLink({
transformer: superjson,
url: getUrl(),
methodOverride: 'POST',
maxItems: 1,
headers: () => {
const headers: Record<string, string> = {};
// Don't set admin viewing user header for playground routes
const isPlaygroundRoute =
typeof window !== 'undefined' && window.location.pathname.includes('/playground');
if (isPlaygroundRoute) {
// Add header to indicate this is from playground route
headers['X-From-Playground'] = 'true';
} else if (adminViewingUserId) {
headers['X-Admin-View-User'] = adminViewingUserId;
}
return headers;
},
// @ts-expect-error - Type conflict between tRPC and the custom fetch wrapper types
fetch: createFetchHandler(),
}),
],
});
// Default tRPC client for direct usage (without admin headers)
// Admin headers are handled by QueryProvider for React Query hooks
export const trpcClient = createTrpcClient(null);
// Export a function to clear the cache (for logout, etc.)
export async function clearQueryCache(connectionId: string | null) {
const effectiveConnectionId = connectionId ?? 'default';
console.log('[Cache Clear] Clearing persisted queries for:', effectiveConnectionId);
await clearPersistedQueries(effectiveConnectionId);
}
// Check if cache restoration is disabled from localStorage
function isCacheRestorationDisabled(): boolean {
if (typeof window === 'undefined') return false;
try {
const value = localStorage.getItem('disableQueryCache');
return value === 'true';
} catch {
return false;
}
}
export function QueryProvider({
children,
connectionId,
}: PropsWithChildren<{ connectionId: string | null }>) {
const adminViewingUserId = useAdminViewingUserId();
const disableQueryCache = isCacheRestorationDisabled();
const isAdminView = adminViewingUserId !== null;
const effectiveConnectionId = connectionId ?? 'default';
const queryClient = useMemo(() => getQueryClient(connectionId), [connectionId]);
const trpcClient = useMemo(() => createTrpcClient(adminViewingUserId), [adminViewingUserId]);
// Track if we've restored the cache for this connection
const [isRestored, setIsRestored] = useState(false);
const restoredConnectionRef = useRef<string | null>(null);
// Restore cached queries on mount (only once per connection)
useEffect(() => {
// Skip in SSR
if (typeof window === 'undefined') return;
// Skip if disabled or admin view
if (disableQueryCache || isAdminView) {
console.log('[QueryProvider] Cache restoration disabled or admin view active');
setIsRestored(true);
return;
}
// Skip if already restored or in-progress for this connection
if (restoredConnectionRef.current === effectiveConnectionId) {
return;
}
// Mark immediately to prevent duplicate restoration if connectionId changes mid-restore
restoredConnectionRef.current = effectiveConnectionId;
// Restore queries from IndexedDB
const restore = async () => {
console.log('[QueryProvider] Restoring queries for connection:', effectiveConnectionId);
const startTime = performance.now();
try {
const { restored, expired, skipped } = await restoreQueries(
queryClient,
effectiveConnectionId,
);
const duration = performance.now() - startTime;
console.log(
`[QueryProvider] Restored ${restored} queries (${expired} expired, ${skipped} newer in cache) in ${duration.toFixed(1)}ms`,
);
} catch (error) {
console.error('[QueryProvider] Failed to restore queries:', error);
}
setIsRestored(true);
};
restore();
}, [effectiveConnectionId, queryClient, disableQueryCache, isAdminView]);
// Setup persistence subscription (write changes to IndexedDB)
useEffect(() => {
// Skip in SSR
if (typeof window === 'undefined') return;
// Skip if disabled or admin view
if (disableQueryCache || isAdminView) {
return;
}
// Wait until restored before setting up persistence
if (!isRestored) {
return;
}
// Setup persistence and return cleanup function
const cleanup = setupQueryPersistence(queryClient, effectiveConnectionId);
return cleanup;
}, [effectiveConnectionId, queryClient, disableQueryCache, isAdminView, isRestored]);
// Holds "a thread marked done stays gone" as an invariant over the
// listThreads cache, and reports every suppression SERVER-side. The report
// goes to the server on purpose: the previous round's diagnostic went to
// PostHog only, and the browser reporting the bug sends PostHog nothing — 66
// archives on staging produced zero client events, so the fix could never be
// verified. See apps/server/src/docs/bug-mark-done-threads-reappear.md.
useEffect(() => {
if (typeof window === 'undefined') return;
setSuppressionReporter((event) => {
// Fire and forget — a diagnostic must never surface an error to the user.
void trpcClient.mail.reportListSuppression.mutate(event).catch(() => {});
});
const uninstall = installArchivedThreadGuard(queryClient);
return () => {
setSuppressionReporter(null);
uninstall();
};
}, [queryClient]);
// Clear cache when admin view changes or on initial mount if admin view is active
useEffect(() => {
if (adminViewingUserId !== null) {
// Clear cache when entering admin view
queryClient.clear();
// Also clear IndexedDB cache to prevent stale data restoration
clearQueryCache(connectionId).catch((error) => {
console.error('[QueryProvider] Error clearing IndexedDB cache on admin view:', error);
});
}
}, [adminViewingUserId, queryClient, connectionId]);
return (
<QueryClientProvider client={queryClient}>
<TRPCProvider trpcClient={trpcClient} queryClient={queryClient}>
{children}
</TRPCProvider>
</QueryClientProvider>
);
}