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. 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.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("-"); }); }); }); 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(); } }