fix: merge sessionChanges before finalizing session after refresh (#2401) (#2414)

Author: tusharpandey13

src/server/client.test.ts

@@ -244,7 +244,13 @@ describe("Auth0Client", () => {
       // Restore mocking of getTokenSet directly
       mockGetTokenSet = vi
         .spyOn(AuthClient.prototype as any, "getTokenSet")
-        .mockResolvedValue([null, mockRefreshedTokenSet]); // Simulate successful refresh
+        .mockResolvedValue([
+          null,
+          {
+            tokenSet: mockRefreshedTokenSet,
+            idTokenClaims: {}
+          }
+        ]); // Simulate successful refresh
 
       // Remove mocks for discoverAuthorizationServerMetadata and getClientAuth
       // Remove fetch mock
@@ -286,6 +292,53 @@ describe("Auth0Client", () => {
       // Verify save was not called
       expect(mockSaveToSession).not.toHaveBeenCalled();
     });
+
+    it("should provide the refreshed accessToken to beforeSessionSaved hook", async () => {
+      let accessToken: string | undefined;
+
+      client = new Auth0Client({
+        beforeSessionSaved: async (session) => {
+          accessToken = session.tokenSet?.accessToken;
+          return session;
+        }
+      });
+
+      const mockReq = { headers: new Headers() } as NextRequest;
+      const mockRes = new NextResponse();
+
+      await client.getAccessToken(mockReq, mockRes, { refresh: true });
+
+      expect(accessToken).toBe("new_access_token");
+    });
+
+    it("should honor changes made to the tokenSet in beforeSessionSaved hook", async () => {
+      client = new Auth0Client({
+        beforeSessionSaved: async (session) => {
+          return {
+            ...session,
+            tokenSet: {
+              ...session.tokenSet,
+              idToken: "modified_id_token"
+            }
+          };
+        }
+      });
+
+      const mockReq = { headers: new Headers() } as NextRequest;
+      const mockRes = new NextResponse();
+
+      await client.getAccessToken(mockReq, mockRes, { refresh: true });
+
+      expect(mockSaveToSession).toHaveBeenCalledWith(
+        expect.objectContaining({
+          tokenSet: expect.objectContaining({
+            idToken: "modified_id_token"
+          })
+        }),
+        mockReq,
+        mockRes
+      );
+    });
   });
 
   describe("constructor configuration", () => {

src/server/client.ts

@@ -695,17 +695,10 @@ export class Auth0Client {
       // call beforeSessionSaved callback if present
       // if not then filter id_token claims with default rules
       const finalSession = await this.authClient.finalizeSession(
-        session,
+        { ...session, ...sessionChanges },
         tokenSet.idToken
       );
-      await this.saveToSession(
-        {
-          ...finalSession,
-          ...sessionChanges
-        },
-        req,
-        res
-      );
+      await this.saveToSession(finalSession, req, res);
     }
 
     return {