feat: add tenant id to lambda context and structured log message

Ramisa Alam · Ramisa Alam · commit b7c13acceca2 · 2025-04-17T10:56:53.000Z
diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java
@@ -581,6 +581,7 @@ public ByteArrayOutputStream call(InvocationRequest request) throws Error, Excep
                         cognitoIdentity,
                         LambdaEnvironment.FUNCTION_VERSION,
                         request.getInvokedFunctionArn(),
+                        request.getTenantId(),
                         clientContext
                 );
 
diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java
@@ -22,6 +22,7 @@ public class LambdaContext implements Context {
     private final long deadlineTimeInMs;
     private final CognitoIdentity cognitoIdentity;
     private final ClientContext clientContext;
+    private final String tenantId;
     private final LambdaLogger logger;
 
     public LambdaContext(
@@ -34,6 +35,7 @@ public LambdaContext(
             CognitoIdentity identity,
             String functionVersion,
             String invokedFunctionArn,
+            String tenantId,
             ClientContext clientContext
     ) {
         this.memoryLimit = memoryLimit;
@@ -46,6 +48,7 @@ public LambdaContext(
         this.clientContext = clientContext;
         this.functionVersion = functionVersion;
         this.invokedFunctionArn = invokedFunctionArn;
+        this.tenantId = tenantId;
         this.logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger();
     }
 
@@ -91,6 +94,10 @@ public int getRemainingTimeInMillis() {
         return delta > 0 ? delta : 0;
     }
 
+    public String getTenantId() {
+        return tenantId;
+    }
+
     public LambdaLogger getLogger() {
         return logger;
     }
diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java
@@ -41,6 +41,7 @@ private StructuredLogMessage createLogMessage(String message, LogLevel logLevel)
 
         if (lambdaContext != null) {
             msg.AWSRequestId = lambdaContext.getAwsRequestId();
+            msg.tenantId = lambdaContext.getTenantId();
         }
         return msg;
     }
diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java
@@ -12,4 +12,5 @@ class StructuredLogMessage {
     public String message;
     public LogLevel level;
     public String AWSRequestId;
+    public String tenantId;
 }
diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java
@@ -40,6 +40,11 @@ public class InvocationRequest {
      */
     private String cognitoIdentity;
 
+    /**
+     * The tenant ID associated with the request.
+     */
+    private String tenantId;
+
     private byte[] content;
 
     public String getId() {
@@ -94,6 +99,14 @@ public void setCognitoIdentity(String cognitoIdentity) {
         this.cognitoIdentity = cognitoIdentity;
     }
 
+    public String getTenantId() {
+        return tenantId;
+    }
+
+    public void setTenantId(String tenantId) {
+        this.tenantId = tenantId;
+    }
+
     public byte[] getContent() {
         return content;
     }
diff --git a/aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp b/aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp
@@ -20,6 +20,7 @@ static jfieldID contentField;
 static jfieldID clientContextField;
 static jfieldID cognitoIdentityField;
 static jfieldID xrayTraceIdField;
+static jfieldID tenantIdField;
 
 
 jint JNI_OnLoad(JavaVM* vm, void* reserved) {
@@ -41,6 +42,7 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) {
     xrayTraceIdField = env->GetFieldID(invocationRequestClass , "xrayTraceId", "Ljava/lang/String;");
     clientContextField = env->GetFieldID(invocationRequestClass , "clientContext", "Ljava/lang/String;");
     cognitoIdentityField = env->GetFieldID(invocationRequestClass , "cognitoIdentity", "Ljava/lang/String;");
+    tenantIdField = env->GetFieldID(invocationRequestClass, "tenantId", "Ljava/lang/String;");
 
     return JNI_VERSION;
 }
@@ -106,6 +108,10 @@ JNIEXPORT jobject JNICALL Java_com_amazonaws_services_lambda_runtime_api_client_
     CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, cognitoIdentityField, env->NewStringUTF(response.cognito_identity.c_str())));
   }
 
+  if(response.tenant_id != ""){
+    CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, tenantIdField, env->NewStringUTF(response.tenant_id.c_str())));
+  }
+
   bytes = reinterpret_cast<const jbyte*>(response.payload.c_str());
   CHECK_EXCEPTION(env, jArray = env->NewByteArray(response.payload.length()));
   CHECK_EXCEPTION(env, env->SetByteArrayRegion(jArray, 0, response.payload.length(), bytes));
diff --git a/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h b/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h
@@ -61,6 +61,11 @@ struct invocation_request {
      */
     std::chrono::time_point<std::chrono::system_clock> deadline;
 
+    /**
+     * Tenant ID of the current invocation.
+     */
+    std::string tenant_id;
+
     /**
      * The number of milliseconds left before lambda terminates the current execution.
      */
diff --git a/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp b/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp
@@ -40,6 +40,7 @@ static constexpr auto CLIENT_CONTEXT_HEADER = "lambda-runtime-client-context";
 static constexpr auto COGNITO_IDENTITY_HEADER = "lambda-runtime-cognito-identity";
 static constexpr auto DEADLINE_MS_HEADER = "lambda-runtime-deadline-ms";
 static constexpr auto FUNCTION_ARN_HEADER = "lambda-runtime-invoked-function-arn";
+static constexpr auto TENANT_ID_HEADER = "lambda-runtime-aws-tenant-id";
 
 enum Endpoints {
     INIT,
@@ -301,6 +302,10 @@ runtime::next_outcome runtime::get_next()
             req.payload.c_str(),
             static_cast<int64_t>(req.get_time_remaining().count()));
     }
+
+    if (resp.has_header(TENANT_ID_HEADER)) {
+        req.tenant_id = resp.get_header(TENANT_ID_HEADER);
+    }
     return next_outcome(req);
 }
 
diff --git a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java
@@ -18,6 +18,7 @@ public class LambdaContextTest {
     private static final String INVOKED_FUNCTION_ARN = "invoked-function-arn";
     private static final LambdaClientContext CLIENT_CONTEXT = new LambdaClientContext();
     public static final int MEMORY_LIMIT = 128;
+    public static final String TENANT_ID = "tenant-id";
 
     @Test
     public void getRemainingTimeInMillis() {
@@ -54,6 +55,6 @@ public void getRemainingTimeInMillis_Deadline() throws InterruptedException {
 
     private LambdaContext createContextWithDeadline(long deadlineTimeInMs) {
         return new LambdaContext(MEMORY_LIMIT, deadlineTimeInMs, REQUEST_ID, LOG_GROUP_NAME, LOG_STREAM_NAME,
-                FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, CLIENT_CONTEXT);
+                FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, TENANT_ID, CLIENT_CONTEXT);
     }
 }
diff --git a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatterTest.java b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatterTest.java
@@ -29,6 +29,25 @@ void testFormattingWithLambdaContext() {
                 null,
                 null,
                 "function-arn",
+                null,
+                null
+        );
+        assertFormatsString("test log", LogLevel.WARN, context);
+    }
+
+    @Test
+    void testFormattingWithTenantIdInLambdaContext() {
+        LambdaContext context = new LambdaContext(
+                0,
+                0,
+                "request-id",
+                null,
+                null,
+                "function-name",
+                null,
+                null,
+                "function-arn",
+                "tenant-id",
                 null
         );
         assertFormatsString("test log", LogLevel.WARN, context);
@@ -52,6 +71,7 @@ void assert_expected_log_message(StructuredLogMessage result, String message, Lo
 
         if (context != null) {
             assertEquals(context.getAwsRequestId(), result.AWSRequestId);
+            assertEquals(context.getTenantId(), result.tenantId);
         }
     }
 }
diff --git a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/LambdaRuntimeApiClientImplTest.java b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/LambdaRuntimeApiClientImplTest.java
@@ -15,6 +15,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.ErrorRequest;
+import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.InvocationRequest;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.StackElement;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayErrorCause;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayException;
@@ -312,27 +313,62 @@ public void restoreNextWrongStatusCodeTest() {
     }
 
     @Test
-    public void nextTest() {
+    public void nextWithoutTenantIdHeaderTest() {
         try {
-            MockResponse mockResponse = new MockResponse();
-            mockResponse.setResponseCode(HTTP_ACCEPTED);
-            mockResponse.setHeader("lambda-runtime-aws-request-id", "1234567890");
-            mockResponse.setHeader("Content-Type", "application/json");
+            MockResponse mockResponse = buildMockResponseForNextInvocation();
             mockWebServer.enqueue(mockResponse);
 
-            lambdaRuntimeApiClientImpl.nextInvocation();
-            RecordedRequest recordedRequest = mockWebServer.takeRequest();
-            HttpUrl actualUrl = recordedRequest.getRequestUrl();
-            String expectedUrl = "http://" + getHostnamePort() + "/2018-06-01/runtime/invocation/next";
-            assertEquals(expectedUrl, actualUrl.toString());
+            InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation();
+            verifyNextInvocationRequest();
+            assertNull(invocationRequest.getTenantId());
+        } catch(Exception e) {
+            fail();
+        }
+    }
+
+    @Test
+    public void nextWithTenantIdHeaderTest() {
+        try {
+            MockResponse mockResponse = buildMockResponseForNextInvocation();
+            String expectedTenantId = "my-tenant-id";
+            mockResponse.setHeader("lambda-runtime-aws-tenant-id", expectedTenantId);
+            mockWebServer.enqueue(mockResponse);
+
+            InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation();
+            verifyNextInvocationRequest();
+            assertEquals(expectedTenantId, invocationRequest.getTenantId());
 
-            String actualBody = recordedRequest.getBody().readUtf8();
-            assertEquals("", actualBody);
         } catch(Exception e) {
             fail();
         }
     }
 
+    @Test
+    public void nextWithEmptyTenantIdHeaderTest() {
+        try {
+            MockResponse mockResponse = buildMockResponseForNextInvocation();
+            mockResponse.setHeader("lambda-runtime-aws-tenant-id", "");
+            mockWebServer.enqueue(mockResponse);
+    
+            InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation();
+            verifyNextInvocationRequest();
+            assertNull(invocationRequest.getTenantId());
+        } catch(Exception e) {
+            fail();
+        }
+    }
+
+    @Test
+    public void nextWithNullTenantIdHeaderTest() {
+        try {
+            MockResponse mockResponse = buildMockResponseForNextInvocation();
+            assertThrows(NullPointerException.class, () -> {
+                mockResponse.setHeader("lambda-runtime-aws-tenant-id", null);
+            });
+        } catch(Exception e) {
+            fail();
+        }
+    }
 
     @Test
     public void createUrlMalformedTest() {
@@ -376,6 +412,24 @@ public void lambdaReportErrorXRayHeaderTooLongTest() {
         }
     }
 
+    private MockResponse buildMockResponseForNextInvocation() {
+        MockResponse mockResponse = new MockResponse();
+        mockResponse.setResponseCode(HTTP_ACCEPTED);
+        mockResponse.setHeader("lambda-runtime-aws-request-id", "1234567890");
+        mockResponse.setHeader("Content-Type", "application/json");
+        return mockResponse;
+    }
+
+    private void verifyNextInvocationRequest() throws Exception {
+        RecordedRequest recordedRequest = mockWebServer.takeRequest();
+        HttpUrl actualUrl = recordedRequest.getRequestUrl();
+        String expectedUrl = "http://" + getHostnamePort() + "/2018-06-01/runtime/invocation/next";
+        assertEquals(expectedUrl, actualUrl.toString());
+
+        String actualBody = recordedRequest.getBody().readUtf8();
+        assertEquals("", actualBody);
+    }
+
     private String getHostnamePort() {
         return mockWebServer.getHostName() + ":" + mockWebServer.getPort();
     }

Original file line number	Diff line number	Diff line change
`@@ -41,6 +41,7 @@ private StructuredLogMessage createLogMessage(String message, LogLevel logLevel)`
`41`	`41`
`42`	`42`	`if (lambdaContext != null) {`
`43`	`43`	`msg.AWSRequestId = lambdaContext.getAwsRequestId();`
	`44`	`+ msg.tenantId = lambdaContext.getTenantId();`
`44`	`45`	`}`
`45`	`46`	`return msg;`
`46`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -12,4 +12,5 @@ class StructuredLogMessage {`
`12`	`12`	`public String message;`
`13`	`13`	`public LogLevel level;`
`14`	`14`	`public String AWSRequestId;`
	`15`	`+ public String tenantId;`
`15`	`16`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@ public class LambdaContextTest {`
`18`	`18`	`private static final String INVOKED_FUNCTION_ARN = "invoked-function-arn";`
`19`	`19`	`private static final LambdaClientContext CLIENT_CONTEXT = new LambdaClientContext();`
`20`	`20`	`public static final int MEMORY_LIMIT = 128;`
	`21`	`+ public static final String TENANT_ID = "tenant-id";`
`21`	`22`
`22`	`23`	`@Test`
`23`	`24`	`public void getRemainingTimeInMillis() {`
`@@ -54,6 +55,6 @@ public void getRemainingTimeInMillis_Deadline() throws InterruptedException {`
`54`	`55`
`55`	`56`	`private LambdaContext createContextWithDeadline(long deadlineTimeInMs) {`
`56`	`57`	`return new LambdaContext(MEMORY_LIMIT, deadlineTimeInMs, REQUEST_ID, LOG_GROUP_NAME, LOG_STREAM_NAME,`
`57`		`- FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, CLIENT_CONTEXT);`
	`58`	`+ FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, TENANT_ID, CLIENT_CONTEXT);`
`58`	`59`	`}`
`59`	`60`	`}`