mosaicmap/lib/apollo-client.ts

import { ApolloClient, InMemoryCache, createHttpLink, from, split } from '@apollo/client';
import { getMainDefinition } from '@apollo/client/utilities';
import { setContext } from '@apollo/client/link/context';
import { GraphQLWsLink } from '@apollo/client/link/subscriptions';
import { createClient, Client } from 'graphql-ws';
// import * as Sentry from '@sentry/nextjs';


const TOKEN_KEY = 'auth_token';

export interface WSConnectionState {
    status: 'connecting' | 'connected' | 'disconnected' | 'reconnecting';
    error?: Error;
    reconnectAttempts: number;
}

let wsConnectionState: WSConnectionState = {
    status: 'disconnected',
    reconnectAttempts: 0
};

const connectionStateListeners: Set<(state: WSConnectionState) => void> = new Set();

export const subscribeToConnectionState = (listener: (state: WSConnectionState) => void) => {
    connectionStateListeners.add(listener);
    listener(wsConnectionState);

    return () => {
        connectionStateListeners.delete(listener);
    };
};

const updateConnectionState = (update: Partial<WSConnectionState>) => {
    wsConnectionState = { ...wsConnectionState, ...update };
    connectionStateListeners.forEach(listener => listener(wsConnectionState));
};

const httpLink = createHttpLink({
    uri: process.env.NEXT_PUBLIC_GRAPHQL_BACKEND_URL || 'http://localhost:3050/graphql',
});

let wsClient: Client | null = null;

const createWSClient = () => {
    const wsUrl = process.env.NEXT_PUBLIC_GRAPHQL_BACKEND_URL?.replace('/graphql', '/ws')?.replace('http://', 'ws://') || 'ws://localhost:3050/ws';

    // console.log('Creating WebSocket client with URL:', wsUrl);

    return createClient({
        url: wsUrl,
        connectionParams: () => {
            const token = typeof window !== 'undefined' ? localStorage.getItem(TOKEN_KEY) : null;
            return token ? { authorization: `Bearer ${token}` } : {};
        },
        retryAttempts: 5,
        retryWait: async function (attempt) {
            const delay = Math.min(1000 * Math.pow(2, attempt), 30000);
            updateConnectionState({
                status: 'reconnecting',
                reconnectAttempts: attempt + 1
            });

            // Sentry.startSpan({
            //     op: 'websocket.reconnect',
            //     name: 'WebSocket Reconnection Attempt'
            // }, (span) => {
            //     span.setAttribute('attempt', attempt + 1);
            //     span.setAttribute('delay', delay);
            // });

            await new Promise(resolve => setTimeout(resolve, delay));
        },
        on: {
            connecting: () => {
                // console.log('WebSocket connecting...');
                updateConnectionState({ status: 'connecting' });
            },
            opened: () => {
                // console.log('WebSocket connection opened successfully');
                updateConnectionState({
                    status: 'connected',
                    error: undefined,
                    reconnectAttempts: 0
                });

                // Sentry.startSpan({
                //     op: 'websocket.connect',
                //     name: 'WebSocket Connection Established'
                // }, () => { });
            },
            closed: (event: any) => {
                // console.log('WebSocket connection closed:', event.code, event.reason);
                const error = event.reason ? new Error(event.reason) : undefined;
                updateConnectionState({
                    status: 'disconnected',
                    error
                });

                // if (error) {
                //     Sentry.captureException(error);
                // }
            },
            error: (error) => {
                console.error('WebSocket error:', error);
                updateConnectionState({
                    status: 'disconnected',
                    error: error instanceof Error ? error : new Error('WebSocket error')
                });

                // Sentry.captureException(error);
            }
        }
    });
};

const getOrCreateWSClient = () => {
    if (!wsClient) {
        wsClient = createWSClient();
    }
    return wsClient;
};

const wsLink = new GraphQLWsLink(getOrCreateWSClient());

const authLink = setContext((_, { headers }) => {
    // 从 localStorage 获取 token
    const token = typeof window !== 'undefined' ? localStorage.getItem(TOKEN_KEY) : null;

    // 如果有 token，添加到 headers 中
    if (token) {
        return {
            headers: {
                ...headers,
                authorization: `Bearer ${token}`,
            }
        };
    }

    // 如果没有 token，返回原始 headers
    return {
        headers: {
            ...headers,
        }
    };
});

const link = split(
    ({ query }) => {
        const def = getMainDefinition(query);
        return def.kind === 'OperationDefinition' && def.operation === 'subscription';
    },
    wsLink,   // 如果是 subscription，则走 ws
    from([authLink, httpLink])
);

export const createApolloClient = () => {
    return new ApolloClient({
        link,
        cache: new InMemoryCache(),
    });
};

export const reconnectWebSocket = () => {
    if (wsClient) {
        wsClient.dispose();
        wsClient = null;
    }
    wsClient = createWSClient();
    return wsClient;
};