From 94e1e85c0b90124b7cd850eae6aa7ffaa481de4c Mon Sep 17 00:00:00 2001 From: Maxime David Date: Fri, 17 Jul 2026 14:28:00 +0000 Subject: [PATCH 1/2] fix: inject tenantId into InvokeStore for custom JSON logs runWithInvokeContext only stored REQUEST_ID and X_RAY_TRACE_ID in the InvokeStore, never TENANT_ID. The JSON loggers read tenantId via invokeStore.getTenantId() (PROTECTED_KEYS.TENANT_ID), so it was always undefined, fell back to "", and formatJsonMessage's `if (tenantId)` guard dropped the field. The header was parsed correctly onto context.tenantId (hence available in code and platform events), but that value never reached the store the loggers consult. Thread context.tenantId through runWithInvokeContext and set PROTECTED_KEYS.TENANT_ID for both the single- and multi-concurrent call sites so tenantId appears in custom JSON log entries, matching the documented behavior and the Python runtime. Add regression tests covering both call sites. tenantId is pinned to the sandbox, so it stays constant across every invocation in the execution environment's lifecycle (unlike requestId/traceId). Ref: V2290595765 cr: https://code.amazon.com/reviews/CR-290016955 --- src/runtime/runtime.test.ts | 72 +++++++++++++++++++++++++++++++++++++ src/runtime/runtime.ts | 4 +++ 2 files changed, 76 insertions(+) diff --git a/src/runtime/runtime.test.ts b/src/runtime/runtime.test.ts index 781cb22..3637fcd 100644 --- a/src/runtime/runtime.test.ts +++ b/src/runtime/runtime.test.ts @@ -353,6 +353,30 @@ describe("Runtime", () => { ]); }); + it("should set tenantId in InvokeStore for buffered invocations", async () => { + // GIVEN + const mockRapidClient = createMockRapidClient(-1, false, true); + const storeOrder: (string | undefined)[] = []; + + const BufferedHandler: BufferedHandler = async () => { + storeOrder.push(invokeStore.getTenantId()); + await new Promise((resolve) => setTimeout(resolve, 50)); // Simulate work + return "done"; + }; + + const runtime = Runtime.create({ + rapidClient: mockRapidClient, + handler: BufferedHandler, + isMultiConcurrent: false, + }); + + // WHEN + await startAndTickRuntime(runtime, 60); + + // THEN - tenantId is pinned to the sandbox, constant across invocations + expect(storeOrder).toEqual(["test-tenant", "test-tenant"]); + }); + it("should set requestId and traceId in InvokeStore for streaming invocations", async () => { // GIVEN const mockRapidClient = createMockRapidClient(3, true); @@ -556,6 +580,45 @@ describe("Runtime", () => { expect(mockRapidClient.postInvocationResponse).toHaveBeenCalledTimes(2); }); + it("should set tenantId in InvokeStore for concurrent invocations", async () => { + // GIVEN + const mockRapidClient = createMockRapidClient(3, false, true); + const storeOrder: (string | undefined)[] = []; + const completionOrder: (string | undefined)[] = []; + + const BufferedHandler: BufferedHandler = async () => { + const invocationNumber = storeOrder.length + 1; + + storeOrder.push(invokeStore.getTenantId()); + + await new Promise((resolve) => + setTimeout(resolve, invocationNumber === 1 ? 10 : 5), + ); + + completionOrder.push(invokeStore.getTenantId()); + + return "done"; + }; + + const runtime = Runtime.create({ + rapidClient: mockRapidClient, + handler: BufferedHandler, + isMultiConcurrent: true, + }); + + // WHEN + runtime.start(); + await new Promise((resolve) => setTimeout(resolve, 20)); + + // THEN - tenantId is pinned to the sandbox, so every concurrent + // invocation reads the same value, both at start and after async gaps. + expect(storeOrder).toEqual(["test-tenant", "test-tenant"]); + expect(completionOrder).toEqual(["test-tenant", "test-tenant"]); + + expect(mockRapidClient.nextInvocation).toHaveBeenCalledTimes(3); + expect(mockRapidClient.postInvocationResponse).toHaveBeenCalledTimes(2); + }); + it("should maintain context when mixing different timer types in RIC", async () => { // GIVEN const mockRapidClient = createMockRapidClient(3, true); @@ -897,6 +960,7 @@ async function startAndTickRuntime( function createMockRapidClient( throttleAt: number = -1, passXRayTraceId: boolean = false, + passTenantId: boolean = false, ): RAPIDClient { let invocationCount = 0; return { @@ -915,6 +979,14 @@ function createMockRapidClient( "lambda-runtime-trace-id": `trace-id-${invocationCount}`, } : {}), + ...(passTenantId + ? { + // tenantId is pinned to the sandbox, so it's constant across + // every invocation in the execution environment's lifecycle + // (unlike requestId/traceId, which change per invocation). + "lambda-runtime-aws-tenant-id": "test-tenant", + } + : {}), }, }); }), diff --git a/src/runtime/runtime.ts b/src/runtime/runtime.ts index 379ee8e..419258c 100644 --- a/src/runtime/runtime.ts +++ b/src/runtime/runtime.ts @@ -67,6 +67,7 @@ export class Runtime { await this.runWithInvokeContext( context.awsRequestId, context.xRayTraceId, + context.tenantId, () => processor.processInvoke(context, event), ); } @@ -82,6 +83,7 @@ export class Runtime { await this.runWithInvokeContext( context.awsRequestId, context.xRayTraceId, + context.tenantId, () => processor.processInvoke(context, event), ); }); @@ -91,6 +93,7 @@ export class Runtime { private async runWithInvokeContext( requestId: string, xRayTraceId: string | undefined, + tenantId: string | undefined, fn: () => Promise, ): Promise { const invokeStore = await InvokeStore.getInstanceAsync(); @@ -98,6 +101,7 @@ export class Runtime { { [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: requestId, [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: xRayTraceId, + [InvokeStoreBase.PROTECTED_KEYS.TENANT_ID]: tenantId, }, fn, ); From 3991a9df71bfb23eb0b9400337671ee7ab4fbab8 Mon Sep 17 00:00:00 2001 From: Maxime David Date: Fri, 17 Jul 2026 18:12:17 +0000 Subject: [PATCH 2/2] test: assert tenantId in JSON log output Reuse the existing JSON-log test to verify tenantId is emitted as a top-level field in JSON-formatted logs, per review feedback on CR-290016955. cr: https://code.amazon.com/reviews/CR-290016955 --- src/logging/log-patch.test.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/logging/log-patch.test.ts b/src/logging/log-patch.test.ts index dfc9c73..7f06bdd 100644 --- a/src/logging/log-patch.test.ts +++ b/src/logging/log-patch.test.ts @@ -10,6 +10,7 @@ describe("LogPatch", async () => { const originalConsole = { ...console }; const mockTimestamp = "2024-01-01T00:00:00.000Z"; const mockRequestId = "test-request-id"; + const mockTenantId = "test-tenant-id"; const invokeStore = await InvokeStore.getInstanceAsync(); @@ -25,6 +26,7 @@ describe("LogPatch", async () => { // Mock InvokeStore vi.spyOn(invokeStore, "getRequestId").mockReturnValue(mockRequestId); + vi.spyOn(invokeStore, "getTenantId").mockReturnValue(mockTenantId); // Mock Date vi.spyOn(Date.prototype, "toISOString").mockReturnValue(mockTimestamp); @@ -230,6 +232,9 @@ describe("LogPatch", async () => { expect(process.stdout.write).toHaveBeenCalledWith( expect.stringContaining(`"requestId":"${mockRequestId}"`), ); + expect(process.stdout.write).toHaveBeenCalledWith( + expect.stringContaining(`"tenantId":"${mockTenantId}"`), + ); expect(process.stdout.write).toHaveBeenCalledWith( expect.stringContaining(`"level":"INFO"`), );