feat: add tenant-id to invoke context

deps/patches/aws-lambda-cpp-add-tenant-id.patch

@@ -0,0 +1,40 @@
+diff --git a/include/aws/lambda-runtime/runtime.h b/include/aws/lambda-runtime/runtime.h
+index f7cb8ef..833b69d 100644
+--- a/include/aws/lambda-runtime/runtime.h
++++ b/include/aws/lambda-runtime/runtime.h
+@@ -56,6 +56,11 @@ struct invocation_request {
+      */
+     std::string function_arn;
+
++    /**
++     * The Tenant ID of the current invocation.
++     */
++    std::string tenant_id;
++
+     /**
+      * Function execution deadline counted in milliseconds since the Unix epoch.
+      */
+diff --git a/src/runtime.cpp b/src/runtime.cpp
+index d1e655f..bcc217d 100644
+--- a/src/runtime.cpp
++++ b/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,
+@@ -296,6 +297,11 @@ runtime::next_outcome runtime::get_next()
+         req.function_arn = std::move(out).get_result();
+     }
+
++    out = resp.get_header(TENANT_ID_HEADER);
++    if (out.is_success()) {
++        req.tenant_id = std::move(out).get_result();
++    }
++
+     out = resp.get_header(DEADLINE_MS_HEADER);
+     if (out.is_success()) {
+         auto const& deadline_string = std::move(out).get_result();

scripts/update_dependencies.sh

@@ -29,7 +29,8 @@ wget -c https://github.com/awslabs/aws-lambda-cpp/archive/refs/tags/v$AWS_LAMBDA
 # Apply patches
 (
   cd aws-lambda-cpp-$AWS_LAMBDA_CPP_RELEASE && \
-    patch -p1 < ../patches/aws-lambda-cpp-add-xray-response.patch
+    patch -p1 < ../patches/aws-lambda-cpp-add-xray-response.patch && \
+    patch -p1 < ../patches/aws-lambda-cpp-add-tenant-id.patch
 )
 
 # Pack again and remove the folder

src/InvokeContext.js

@@ -18,6 +18,7 @@ const INVOKE_HEADER = {
   AWSRequestId: 'lambda-runtime-aws-request-id',
   DeadlineMs: 'lambda-runtime-deadline-ms',
   XRayTrace: 'lambda-runtime-trace-id',
+  TenantId: 'lambda-runtime-aws-tenant-id',
 };
 
 module.exports = class InvokeContext {
@@ -91,6 +92,7 @@ module.exports = class InvokeContext {
       ),
       invokedFunctionArn: this.headers[INVOKE_HEADER.ARN],
       awsRequestId: this.headers[INVOKE_HEADER.AWSRequestId],
+      tenantId: this.headers[INVOKE_HEADER.TenantId],
       getRemainingTimeInMillis: function () {
         return deadline - Date.now();
       },

src/rapid-client.cc

@@ -82,6 +82,11 @@ class RuntimeApiNextPromiseWorker : public Napi::AsyncWorker {
                     Napi::String::New(env, "lambda-runtime-cognito-identity"),
                     Napi::String::New(env, response.cognito_identity.c_str()));
         }
+        if (response.tenant_id != "") {
+            headers.Set(
+                    Napi::String::New(env, "lambda-runtime-aws-tenant-id"),
+                    Napi::String::New(env, response.tenant_id.c_str()));
+        }
 
         auto ret = Napi::Object::New(env);
         ret.Set(Napi::String::New(env, "bodyJson"), response_data);

test/unit/InvokeContextTest.js

@@ -35,3 +35,32 @@ describe('Getting remaining invoke time', () => {
     remainingTime.should.lessThanOrEqual(1000);
   });
 });
+
+describe('Verifying tenant id', () => {
+  it('should return undefined if tenant id is not set', () => {
+    let ctx = new InvokeContext({});
+
+    (ctx._headerData().tenantId === undefined).should.be.true();
+  });
+  it('should return undefined if tenant id is set to undefined', () => {
+    let ctx = new InvokeContext({
+      'lambda-runtime-aws-tenant-id': undefined,
+    });
+
+    (ctx._headerData().tenantId === undefined).should.be.true();
+  });
+  it('should return empty if tenant id is set to empty string', () => {
+    let ctx = new InvokeContext({
+      'lambda-runtime-aws-tenant-id': '',
+    });
+
+    (ctx._headerData().tenantId === '').should.be.true();
+  });
+  it('should return the same id if a valid tenant id is set', () => {
+    let ctx = new InvokeContext({
+      'lambda-runtime-aws-tenant-id': 'blue',
+    });
+
+    ctx._headerData().tenantId.should.equal('blue');
+  });
+});

test/unit/RAPIDClientTest.js

@@ -47,6 +47,23 @@ class NoOpNativeHttp {
   }
 }
 
+class MockNativeClient {
+  constructor(response) {
+    this.response = response;
+    this.called = false;
+    this.shouldThrowError = false;
+  }
+
+  next() {
+    this.called = true;
+    if (this.shouldThrowError) {
+      return Promise.reject(new Error('Failed to get next invocation'));
+    } else {
+      return Promise.resolve(this.response);
+    }
+  }
+}
+
 class EvilError extends Error {
   get name() {
     throw 'gotcha';
@@ -115,3 +132,82 @@ describe('invalid request id works', () => {
     });
   });
 });
+
+describe('next invocation with native client works', () => {
+  it('should call the native client next() method', async () => {
+    const mockNative = new MockNativeClient({
+      bodyJson: '',
+      headers: {
+        'lambda-runtime-aws-request-id': 'test-request-id',
+      },
+    });
+    const client = new RAPIDClient('notUsed:1337', undefined, mockNative);
+    client.useAlternativeClient = false;
+
+    await client.nextInvocation();
+    // verify native client was called
+    mockNative.called.should.be.true();
+  });
+  it('should parse all required headers', async () => {
+    const mockResponse = {
+      bodyJson: '{"message":"Hello from Lambda!"}',
+      headers: {
+        'lambda-runtime-aws-request-id': 'test-request-id',
+        'lambda-runtime-deadline-ms': 1619712000000,
+        'lambda-runtime-trace-id': 'test-trace-id',
+        'lambda-runtime-invoked-function-arn': 'test-function-arn',
+        'lambda-runtime-client-context': '{"client":{"app_title":"MyApp"}}',
+        'lambda-runtime-cognito-identity':
+          '{"identityId":"id123","identityPoolId":"pool123"}',
+        'lambda-runtime-aws-tenant-id': 'test-tenant-id',
+      },
+    };
+
+    const mockNative = new MockNativeClient(mockResponse);
+    const client = new RAPIDClient('notUsed:1337', undefined, mockNative);
+
+    client.useAlternativeClient = false;
+    const response = await client.nextInvocation();
+
+    // Verify all headers are present
+    response.headers.should.have.property(
+      'lambda-runtime-aws-request-id',
+      'test-request-id',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-deadline-ms',
+      1619712000000,
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-trace-id',
+      'test-trace-id',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-invoked-function-arn',
+      'test-function-arn',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-client-context',
+      '{"client":{"app_title":"MyApp"}}',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-cognito-identity',
+      '{"identityId":"id123","identityPoolId":"pool123"}',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-aws-tenant-id',
+      'test-tenant-id',
+    );
+    // Verify body is correctly passed through
+    response.bodyJson.should.equal('{"message":"Hello from Lambda!"}');
+  });
+  it('should handle native client errors', async () => {
+    const nativeClient = new MockNativeClient({});
+    nativeClient.shouldThrowError = true;
+
+    const client = new RAPIDClient('localhost:8080', null, nativeClient);
+    client.useAlternativeClient = false;
+
+    await client.nextInvocation().should.be.rejectedWith('Failed to get next invocation');
+  });
+});