Apollo Websocket无法使用最新JWT令牌初始化新连接问题咨询

我有一个通过vue-apollo连接GraphQL服务器的Vue应用，对应的vue-apollo.js配置代码如下：

const httpEndpoint = process.env.XXX;
const wsEndpoint = process.env.XXX;

let baseLink = new ApolloLink((operation, forward) => { 
    let isTokenExpired = isJwtExpired(token)
    console.log('isExpired is:',isTokenExpired) ;
    if (isTokenExpired) {
        Auth.currentSession().then((data) => {
            token = data.idToken.jwtToken
            localStorage.setItem('JoT', JSON.stringify(token));
            const headerData = token? {authorization: `Bearer ${token}`,}: {};
            operation.setContext({headers: headerData})
            return forward(operation)
        })
    } else { 
        token = JSON.parse(token);
        const headerData = token? {authorization: `Bearer ${token}`,}: {};
        operation.setContext({headers: headerData})
        return forward(operation)
    }
    
})
let observerLink = new ApolloLink((operation, forward) => { 
    return forward(operation)
})
const baseAndObserverLink = baseLink.concat(observerLink)

const errorLink = onError(error => {
    console.log(JSON.stringify(error))
        if (error.graphQLErrors) {
            console.log("GraphQL Error detected")
            if (error.graphQLErrors.extensions && error.graphQLErrors.extensions.code && error.graphQLErrors.extensions.code === 'invalid-jwt') {
                console.log("JWT EXPIRED")
            }
        } else if (error.networkError && error.networkError.extensions && error.networkError.extensions.code && error.networkError.extensions.code === 'start-failed') { 
            console.log("GraphQL Error detected type 2")
                console.log("Unable to Connect")
            
        }else if (error.networkError && error.networkError.extensions && error.networkError.extensions.code && error.networkError.extensions.code === 'validation-failed') { 
                console.log("GraphQL Error detected type 3")
                console.log("Validation Error")
        }
        else { 
            console.log(JSON.stringify(error))
        }
})
var wsLink = new WebSocketLink({
    uri: wsEndpoint,
    options: {
        reconnect: true,
        lazy: true,
        timeout: 30000,
        connectionParams: async () => {
            const token = JSON.parse(localStorage.getItem('JWT'));
            return {
                headers: {
                Authorization: token ? `Bearer ${token}` : "",
                },
            }
        },
    },
});
 


export var filesRoot =
    process.env.VUE_APP_FILES_ROOT ||
    httpEndpoint.substr(0, httpEndpoint.indexOf('/graphql'));

Vue.prototype.$filesRoot = filesRoot;

var httpLink = new HttpLink({
    uri: httpEndpoint,
});

var splitLLink = split(
    ({ query }) => {
        const definition = getMainDefinition(query);
        return (
            definition.kind === 'OperationDefinition' &&
            definition.operation === 'subscription'
        );
    },
    wsLink,
    httpLink
);

export var defaultOptions = {
    httpEndpoint,
    wsEndpoint,
    tokenName: 'JWT',
    persisting: false,
    websocketsOnly: true,
    ssr: false,
    link:baseAndObserverLink.concat(errorLink).concat(splitLLink),

};

export function createProvider(options = {}) {
    const { apolloClient, wsClient } = createApolloClient({
        ...defaultOptions,
        ...options,
    });
    apolloClient.wsClient = wsClient;
    const apolloProvider = new VueApollo({
        defaultClient: apolloClient,
        defaultOptions: {
            $query: {
                // fetchPolicy: 'cache-and-network',
            },
        },
        async errorHandler(error) {

        },
    });
    return apolloProvider;
}

现有代码首次登录运行正常，重新获取新JWT令牌后，HTTP查询链路能正常携带新令牌，查询功能无异常，但WebSocket连接重连时始终使用旧令牌发起连接，只能通过强制刷新页面重载应用，让WSS初始化时读取最新令牌才能恢复，需要无需重载页面的解决方案。

• 存在存储key不一致的bug：刷新token后写入localStorage的key是JoT，但WebSocket连接的connectionParams读取的key是JWT，哪怕重连也读不到新token。
• ApolloLink里的token刷新逻辑是异步的，没有返回Promise，会导致请求提前发送，上下文更新不生效。
• 核心原因：WebSocket连接建立后不会自动感知token更新，自动重连逻辑默认只在网络断开等异常场景触发，token刷新后不会主动断开旧连接重连，所以一直沿用旧连接的认证信息。
• 额外问题：token序列化逻辑混乱，反复JSON.parse/JSON.stringify容易出现取值错误。

按以下步骤修改即可，不需要刷新页面：

1. 统一token存储逻辑，修正key不一致问题
   去掉不必要的JSON序列化操作，所有读写localStorage的key统一为JWT，初始化时直接从本地缓存取token：

   // 初始化直接读取本地token，避免反复parse
   let token = localStorage.getItem('JWT')
   

2. 修正baseLink的异步逻辑，token刷新成功后主动触发WebSocket重连
   注意异步分支必须return Promise，否则Apollo Link不会等待token更新完成就转发请求；拿到新token后主动关闭旧WS连接，触发重连，重连时connectionParams会自动读取最新的本地token（可以把baseLink的定义挪到createProvider内部，通过闭包持有apolloClient/ wsClient引用）：

   let baseLink = new ApolloLink((operation, forward) => { 
       let isTokenExpired = isJwtExpired(token)
       console.log('isExpired is:',isTokenExpired) ;
       if (isTokenExpired) {
           // 必须return这个Promise，否则请求会提前发送
           return Auth.currentSession().then((data) => {
               token = data.idToken.jwtToken
               // 统一存储key为JWT，不需要JSON.stringify
               localStorage.setItem('JWT', token);
               const headerData = token ? {authorization: `Bearer ${token}`,} : {};
               operation.setContext({headers: headerData})
               // 主动触发WS重连
               if (apolloClient.wsClient) {
                   apolloClient.wsClient.close()
                   apolloClient.wsClient.connect()
               }
               return forward(operation)
           })
       } else { 
           const headerData = token ? {authorization: `Bearer ${token}`,} : {};
           operation.setContext({headers: headerData})
           return forward(operation)
       }
   })
   

3. 同步修改WS连接的connectionParams逻辑，去掉不必要的JSON.parse：

   var wsLink = new WebSocketLink({
       uri: wsEndpoint,
       options: {
           reconnect: true,
           lazy: true,
           timeout: 30000,
           connectionParams: async () => {
               // 直接读最新的token，不需要parse
               const token = localStorage.getItem('JWT');
               return {
                   headers: {
                       Authorization: token ? `Bearer ${token}` : "",
                   },
               }
           },
       },
   });
   

4. 补充errorLink的invalid-jwt处理逻辑，捕获到WS连接返回的token失效错误时，主动触发一次token刷新和重连，避免订阅一直卡在错误状态。

注意：如果使用的是graphql-ws库替代旧的subscriptions-transport-ws，主动重连的方法需要替换为wsClient.dispose()后重新创建客户端实例，从当前配置参数来看用的是旧版库，close()+connect()的调用方式可以生效。

内容的提问来源于stack exchange，提问作者sachsom