fix: include tenantId in custom JSON log entries

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"`),
       );

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",
+              }
+            : {}),
         },
       });
     }),

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,13 +93,15 @@ export class Runtime {
   private async runWithInvokeContext(
     requestId: string,
     xRayTraceId: string | undefined,
+    tenantId: string | undefined,
     fn: () => Promise<void>,
   ): Promise<void> {
     const invokeStore = await InvokeStore.getInstanceAsync();
     return invokeStore.run(
       {
         [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: requestId,
         [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: xRayTraceId,
+        [InvokeStoreBase.PROTECTED_KEYS.TENANT_ID]: tenantId,
       },
       fn,
     );