From 0506e1e98ed67abf749d2769729ff235d089337d Mon Sep 17 00:00:00 2001
From: Wolfgang Goedel <wgoedel@gmail.com>
Date: Mon, 20 Oct 2025 16:43:33 +0200
Subject: [PATCH] merge sessionChanges before finalizing session after refresh

---
 src/server/client.test.ts | 55 ++++++++++++++++++++++++++++++++++++++-
 src/server/client.ts      | 11 ++------
 2 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/src/server/client.test.ts b/src/server/client.test.ts
index ea590154..b58e635b 100644
--- a/src/server/client.test.ts
+++ b/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", () => {
diff --git a/src/server/client.ts b/src/server/client.ts
index 5424538f..b97866d5 100644
--- a/src/server/client.ts
+++ b/src/server/client.ts
@@ -704,17 +704,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 {