From 15beadd1d1e62861abd3c6f33d9ad0ddd6659c9a Mon Sep 17 00:00:00 2001 From: Naman Trivedi Date: Wed, 3 Dec 2025 14:32:05 +0000 Subject: [PATCH 1/3] bug: fix context cleared prematurely in InvokeStoreSingle with async functions --- src/invoke-store.async-context.spec.ts | 27 ++++++++++++++++++++++++++ src/invoke-store.ts | 6 +----- 2 files changed, 28 insertions(+), 5 deletions(-) create mode 100644 src/invoke-store.async-context.spec.ts diff --git a/src/invoke-store.async-context.spec.ts b/src/invoke-store.async-context.spec.ts new file mode 100644 index 0000000..c60d20e --- /dev/null +++ b/src/invoke-store.async-context.spec.ts @@ -0,0 +1,27 @@ +import { describe, it, expect, beforeEach } from "vitest"; +import { InvokeStore, InvokeStoreBase } from "./invoke-store"; + +describe("InvokeStore - Async Context Bug", () => { + let invokeStore: InvokeStoreBase; + + beforeEach(async () => { + if (InvokeStore._testing) { + InvokeStore._testing.reset(); + } + invokeStore = await InvokeStore.getInstanceAsync(); + }); + + it("should not clear context after await in InvokeStoreSingle", async () => { + const testContext = { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-123", + }; + + await invokeStore.run(testContext, async () => { + expect(invokeStore.getRequestId()).toBe("test-123"); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + expect(invokeStore.getRequestId()).toBe("test-123"); + }); + }); +}); diff --git a/src/invoke-store.ts b/src/invoke-store.ts index cc3d3a5..5421974 100644 --- a/src/invoke-store.ts +++ b/src/invoke-store.ts @@ -92,11 +92,7 @@ class InvokeStoreSingle extends InvokeStoreBase { run(context: Context, fn: () => T): T { this.currentContext = context; - try { - return fn(); - } finally { - this.currentContext = undefined; - } + return fn(); } } From 9bed56d892dbe8473a9b8b7f3c4aa5f4e6612057 Mon Sep 17 00:00:00 2001 From: Naman Trivedi Date: Wed, 3 Dec 2025 15:12:44 +0000 Subject: [PATCH 2/3] chore: add changeset for async context fix --- .changeset/fix-async-context.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/fix-async-context.md diff --git a/.changeset/fix-async-context.md b/.changeset/fix-async-context.md new file mode 100644 index 0000000..ac69546 --- /dev/null +++ b/.changeset/fix-async-context.md @@ -0,0 +1,5 @@ +--- +"@aws/lambda-invoke-store": patch +--- + +Fix context cleared prematurely in InvokeStoreSingle with async functions. Removed try-finally block that was clearing context before async operations completed. From 5aee9a2ddda4db7559f70c7c0f9ee13b6c5b01e7 Mon Sep 17 00:00:00 2001 From: Naman Trivedi Date: Wed, 3 Dec 2025 16:01:46 +0000 Subject: [PATCH 3/3] test: update tests to reflect context persistence behavior --- src/invoke-store.spec.ts | 51 ++++++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/src/invoke-store.spec.ts b/src/invoke-store.spec.ts index eff2e6d..be2c674 100644 --- a/src/invoke-store.spec.ts +++ b/src/invoke-store.spec.ts @@ -21,15 +21,6 @@ describe.each([ invokeStore = await InvokeStore.getInstanceAsync(); describe("getRequestId and getXRayTraceId", () => { - it("should return placeholder when called outside run context", () => { - // WHEN - const requestId = invokeStore.getRequestId(); - const traceId = invokeStore.getXRayTraceId(); - - // THEN - expect(requestId).toBe("-"); - expect(traceId).toBeUndefined(); - }); it("should return current invoke IDs when called within run context", async () => { // WHEN @@ -100,12 +91,31 @@ describe.each([ }); describe("getContext", () => { - it("should return undefined when outside run context", () => { - // WHEN - const context = invokeStore.getContext(); + it("should replace context on subsequent run calls", async () => { + // WHEN - First run + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "first-id", + }, + () => { + expect(invokeStore.getRequestId()).toBe("first-id"); + }, + ); - // THEN - expect(context).toBeUndefined(); + // WHEN - Second run should replace context + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "second-id", + }, + () => { + // THEN - Should have new context, not old one + expect(invokeStore.getRequestId()).toBe("second-id"); + const context = invokeStore.getContext(); + expect(context).toEqual({ + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "second-id", + }); + }, + ); }); it("should return complete context with Lambda and custom fields", async () => { @@ -133,14 +143,6 @@ describe.each([ }); describe("hasContext", () => { - it("should return false when outside run context", () => { - // WHEN - const hasContext = invokeStore.hasContext(); - - // THEN - expect(hasContext).toBe(false); - }); - it("should return true when inside run context", async () => { // WHEN const result = await invokeStore.run( @@ -158,7 +160,7 @@ describe.each([ }); describe("error handling", () => { - it("should propagate errors while maintaining isolation", async () => { + it("should propagate errors", async () => { // GIVEN const error = new Error("test error"); @@ -174,7 +176,6 @@ describe.each([ // THEN await expect(promise).rejects.toThrow(error); - expect(invokeStore.getRequestId()).toBe("-"); }); it("should handle errors in concurrent executions independently", async () => { @@ -205,7 +206,6 @@ describe.each([ // THEN expect(traces).toContain("success-success-id"); expect(traces).toContain("before-error-error-id"); - expect(invokeStore.getRequestId()).toBe("-"); }); }); @@ -242,7 +242,6 @@ describe.each([ // THEN await expect(promise).rejects.toThrow(error); - expect(invokeStore.getRequestId()).toBe("-"); }); }); });