From b1552f68546557a45f3a39cf5f801578ef67eef1 Mon Sep 17 00:00:00 2001
From: Hui Zhao <zhohz@amazon.com>
Date: Tue, 21 Nov 2023 15:55:11 -0800
Subject: [PATCH] fix(aws-amplify): destroyAmplifyServerContext may not be
 called

---
 .../runWithAmplifyServerContext.test.ts       | 29 ++++++++++++++++++-
 .../runWithAmplifyServerContext.ts            | 11 ++++---
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/packages/aws-amplify/__tests__/adapterCore/runWithAmplifyServerContext.test.ts b/packages/aws-amplify/__tests__/adapterCore/runWithAmplifyServerContext.test.ts
index 40ccbe1c5c3..1c9c727c627 100644
--- a/packages/aws-amplify/__tests__/adapterCore/runWithAmplifyServerContext.test.ts
+++ b/packages/aws-amplify/__tests__/adapterCore/runWithAmplifyServerContext.test.ts
@@ -29,6 +29,10 @@ describe('runWithAmplifyServerContext', () => {
 		mockCreateAmplifyServerContext.mockReturnValueOnce(mockContextSpec);
 	});
 
+	afterEach(() => {
+		mockDestroyAmplifyServerContext.mockReset();
+	});
+
 	it('should run the operation with the context', () => {
 		const mockOperation = jest.fn();
 		runWithAmplifyServerContext(
@@ -45,7 +49,7 @@ describe('runWithAmplifyServerContext', () => {
 		expect(mockOperation).toHaveBeenCalledWith(mockContextSpec);
 	});
 
-	it('should destroy the context after the operation', async () => {
+	it('should destroy the context after the operation completed', async () => {
 		const mockOperation = jest.fn();
 		await runWithAmplifyServerContext(
 			mockAmplifyConfig,
@@ -63,6 +67,29 @@ describe('runWithAmplifyServerContext', () => {
 		);
 	});
 
+	it('should destroy the context when the operation throws', async () => {
+		const testError = new Error('some error');
+		const mockOperation = jest.fn();
+		mockOperation.mockRejectedValueOnce(testError);
+
+		await expect(
+			runWithAmplifyServerContext(
+				mockAmplifyConfig,
+				{
+					Auth: {
+						tokenProvider: mockTokenProvider,
+						credentialsProvider: mockCredentialAndIdentityProvider,
+					},
+				},
+				mockOperation
+			)
+		).rejects.toThrow(testError);
+
+		expect(mockDestroyAmplifyServerContext).toHaveBeenCalledWith(
+			mockContextSpec
+		);
+	});
+
 	it('should return the result returned by the operation callback function', async () => {
 		const mockResultValue = {
 			url: 'http://123.com',
diff --git a/packages/aws-amplify/src/adapterCore/runWithAmplifyServerContext.ts b/packages/aws-amplify/src/adapterCore/runWithAmplifyServerContext.ts
index 237042f5056..021bfbe0aac 100644
--- a/packages/aws-amplify/src/adapterCore/runWithAmplifyServerContext.ts
+++ b/packages/aws-amplify/src/adapterCore/runWithAmplifyServerContext.ts
@@ -26,9 +26,12 @@ export const runWithAmplifyServerContext: AmplifyServer.RunOperationWithContext
 		);
 
 		// run the operation with injecting the context
-		const result = await operation(contextSpec);
+		try {
+			const result = await operation(contextSpec);
 
-		destroyAmplifyServerContext(contextSpec);
-
-		return result;
+			return result;
+		} finally {
+			// ensures destroy the context regardless whether the operation succeeded or failed
+			destroyAmplifyServerContext(contextSpec);
+		}
 	};