fix: Local Server should pass HTTP headers down to the Lambda Runtime (#643)

<!--- Provide a general summary of your changes in the Title above -->

## Issue \#
<!--- If it fixes an issue, please link to the issue here -->
https://github.com/awslabs/swift-aws-lambda-runtime/issues/607

## Description of changes
<!--- Why is this change required? What problem does it solve? -->
The local HTTP server was not forwarding user‑provided headers to the
runtime’s response. It passes all headers through to the runtime. This
it makes local behavior match the Lambda runtime API contract and allows
developers to opt into metadata by sending the appropriate runtime
headers.

## New/existing dependencies impact assessment, if applicable
<!--- No new dependencies were added to this change. -->
<!--- If any dependency was added / modified / removed,
THIRD-PARTY-LICENSES must be updated accordingly. -->
N/A

## Conventional Commits
<!--- Please use conventional commits to let us know what kind of change
this is.-->
<!--- More info can be found here:
https://www.conventionalcommits.org/en/v1.0.0/-->

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.

---------

Co-authored-by: Sébastien Stormacq <sebastien.stormacq@gmail.com>

Examples/MultiTenant/Package.swift

@@ -6,7 +6,8 @@ let package = Package(
     name: "swift-aws-lambda-runtime-example",
     platforms: [.macOS(.v15)],
     products: [
-        .executable(name: "MultiTenant", targets: ["MultiTenant"])
+        .executable(name: "MultiTenant", targets: ["MultiTenant"]),
+        .executable(name: "MultiTenantLocal", targets: ["MultiTenantLocal"]),
     ],
     dependencies: [
         // For local development (default)
@@ -24,6 +25,12 @@ let package = Package(
                 .product(name: "AWSLambdaRuntime", package: "swift-aws-lambda-runtime"),
                 .product(name: "AWSLambdaEvents", package: "swift-aws-lambda-events"),
             ]
-        )
+        ),
+        .executableTarget(
+            name: "MultiTenantLocal",
+            dependencies: [
+                .product(name: "AWSLambdaRuntime", package: "swift-aws-lambda-runtime")
+            ]
+        ),
     ]
 )

Examples/MultiTenant/README.md

@@ -6,7 +6,8 @@ This example demonstrates how to build a multi-tenant Lambda function using Swif
 
 This example implements a request tracking system that maintains separate counters and request histories for each tenant. The Lambda function:
 
-- Accepts requests from multiple tenants via API Gateway
+- Accepts requests from multiple tenants via API Gateway (MultiTenant)
+- Accepts direct JSON payloads without API Gateway (MultiTenantLocal)
 - Maintains isolated execution environments per tenant
 - Tracks request counts and timestamps for each tenant
 - Returns tenant-specific data in JSON format
@@ -41,6 +42,8 @@ The example consists of:
 
 3. **Lambda Handler** - Processes API Gateway requests and manages tenant data
 
+4. **MultiTenantLocal Handler** - Processes plain JSON requests without API Gateway
+
 ## Code Structure
 
 ```swift
@@ -175,6 +178,24 @@ requestParameters:
 
 ## Testing
 
+### Local Testing (No API Gateway)
+
+Run the simple handler locally and invoke it directly using the local server. This example does not depend on API Gateway or its request/response types.
+
+1. Start the local server:
+   ```bash
+   swift run MultiTenantLocal
+   ```
+
+2. Invoke with a tenant header and a JSON body:
+   ```bash
+   # Test multi-tenant function locally
+   curl -X POST http://127.0.0.1:7000/invoke \
+     -H "Content-Type: application/json" \
+     -H "Lambda-Runtime-Aws-Tenant-Id: tenant-123" \
+     -d '{"message" :  "hello"}'
+   ```
+
 ### Using API Gateway
 
 The tenant ID is passed as a query parameter. API Gateway automatically maps it to the `X-Amz-Tenant-Id` header:

Examples/MultiTenant/Sources/MultiTenantLocal/main.swift

@@ -0,0 +1,117 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftAWSLambdaRuntime open source project
+//
+// Copyright SwiftAWSLambdaRuntime project authors
+// Copyright (c) Amazon.com, Inc. or its affiliates.
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import AWSLambdaRuntime
+
+#if canImport(FoundationEssentials)
+import FoundationEssentials
+#else
+import Foundation
+#endif
+
+let tenants = TenantDataStore()
+
+struct TenantMessage: Codable {
+    let message: String
+}
+
+enum TenantError: Error {
+    case missingTenantID
+}
+
+struct SimpleTenantHandler: LambdaHandler {
+    typealias Event = TenantMessage
+    typealias Output = TenantData
+
+    let tenants: TenantDataStore
+
+    func handle(_ event: TenantMessage, context: LambdaContext) async throws -> TenantData {
+        guard let tenantID = context.tenantID else {
+            throw TenantError.missingTenantID
+        }
+
+        let currentData = await tenants[tenantID] ?? TenantData(tenantID: tenantID)
+        let updatedData = currentData.addingRequest(message: event.message)
+        await tenants.update(id: tenantID, data: updatedData)
+        return updatedData
+    }
+}
+
+let handler = SimpleTenantHandler(tenants: tenants)
+let runtime = LambdaRuntime(handler: LambdaCodableAdapter(handler: LambdaHandlerAdapter(handler: handler)))
+
+try await runtime.run()
+
+actor TenantDataStore {
+    private var tenants: [String: TenantData] = [:]
+
+    subscript(id: String) -> TenantData? {
+        tenants[id]
+    }
+
+    // subscript setters can't be called from outside of the actor
+    func update(id: String, data: TenantData) {
+        tenants[id] = data
+    }
+}
+
+struct TenantData: Codable {
+    struct TenantRequest: Codable {
+        let requestNumber: Int
+        let timestamp: String
+        let message: String
+    }
+
+    let tenantID: String
+    let requestCount: Int
+    let firstRequest: String
+    let requests: [TenantRequest]
+
+    init(tenantID: String) {
+        self.init(
+            tenantID: tenantID,
+            requestCount: 0,
+            firstRequest: "\(Date().timeIntervalSince1970)",
+            requests: []
+        )
+    }
+
+    func addingRequest(message: String) -> TenantData {
+        let newCount = requestCount + 1
+        let newRequest = TenantRequest(
+            requestNumber: newCount,
+            timestamp: "\(Date().timeIntervalSince1970)",
+            message: message
+        )
+        return TenantData(
+            tenantID: tenantID,
+            requestCount: newCount,
+            firstRequest: firstRequest,
+            requests: requests + [newRequest]
+        )
+    }
+
+    private init(
+        tenantID: String,
+        requestCount: Int,
+        firstRequest: String,
+        requests: [TenantRequest]
+    ) {
+        self.tenantID = tenantID
+        self.requestCount = requestCount
+        self.firstRequest = firstRequest
+        self.requests = requests
+    }
+}

Sources/AWSLambdaRuntime/HTTPServer/Lambda+LocalServer.swift

@@ -383,7 +383,9 @@ internal struct LambdaHTTPServer {
             logger[metadataKey: "requestId"] = "\(requestId)"
 
             logger.trace("/invoke received invocation, pushing it to the pool and wait for a lambda response")
-            self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body))
+            self.invocationPool.push(
+                LocalServerInvocation(requestId: requestId, request: body, headers: head.headers)
+            )
 
             // wait for the lambda function to process the request
             // Handle streaming responses by collecting all chunks for this requestId
@@ -518,7 +520,10 @@ internal struct LambdaHTTPServer {
 
         var headers = response.headers ?? HTTPHeaders()
         if let body = response.body {
-            headers.add(name: "Content-Length", value: "\(body.readableBytes)")
+            // Avoid adding Content-Length if already provided (e.g. forwarded from /invoke)
+            if !headers.contains(name: "Content-Length") && !headers.contains(name: "Transfer-Encoding") {
+                headers.add(name: "Content-Length", value: "\(body.readableBytes)")
+            }
         }
 
         if let status = response.status {
@@ -569,11 +574,12 @@ internal struct LambdaHTTPServer {
     struct LocalServerInvocation: Sendable {
         let requestId: String
         let request: ByteBuffer
+        let headers: HTTPHeaders
 
         func acceptedResponse() -> LocalServerResponse {
 
             // required headers
-            let headers = HTTPHeaders([
+            var headers = HTTPHeaders([
                 (AmazonHeaders.requestID, self.requestId),
                 (
                     AmazonHeaders.invokedFunctionARN,
@@ -583,6 +589,9 @@ internal struct LambdaHTTPServer {
                 (AmazonHeaders.deadline, "\(LambdaClock.maxLambdaDeadline)"),
             ])
 
+            // Forward all invocation headers
+            headers.add(contentsOf: self.headers)
+
             return LocalServerResponse(
                 id: self.requestId,
                 status: .accepted,

Tests/AWSLambdaRuntimeTests/LambdaLocalServerTests.swift

@@ -181,11 +181,62 @@ struct LambdaLocalServerTest {
         }
     }
 
-    private func makeInvokeRequest(host: String, port: Int, payload: String) async throws -> (Data, HTTPURLResponse) {
+    @Test("Local server forwards invocation headers to LambdaContext")
+    @available(LambdaSwift 2.0, *)
+    func testLocalServerForwardsInvocationHeaders() async throws {
+        let customPort = 8082
+
+        setenv("LOCAL_LAMBDA_PORT", "\(customPort)", 1)
+        defer { unsetenv("LOCAL_LAMBDA_PORT") }
+
+        let logger = Logger(label: "test", factory: { _ in SwiftLogNoOpLogHandler() })
+
+        let tenantId = try await withTimeout(deadline: .seconds(5)) {
+            async let serverTask: String? = LambdaHTTPServer.withLocalServer(
+                host: "127.0.0.1",
+                port: customPort,
+                invocationEndpoint: nil,
+                logger: logger
+            ) {
+                try await LambdaRuntimeClient.withRuntimeClient(
+                    configuration: .init(ip: "127.0.0.1", port: customPort),
+                    eventLoop: Lambda.defaultEventLoop,
+                    logger: logger
+                ) { runtimeClient in
+                    let (invocation, writer) = try await runtimeClient.nextInvocation()
+                    try await writer.writeAndFinish(ByteBuffer(string: "\"ok\""))
+                    return invocation.metadata.tenantID
+                }
+            }
+
+            try await Task.sleep(for: .milliseconds(200))
+
+            _ = try await self.makeInvokeRequest(
+                host: "127.0.0.1",
+                port: customPort,
+                payload: "\"ping\"",
+                headers: ["Lambda-Runtime-Aws-Tenant-Id": "123"]
+            )
+
+            return try await serverTask
+        }
+
+        #expect(tenantId == "123")
+    }
+
+    private func makeInvokeRequest(
+        host: String,
+        port: Int,
+        payload: String,
+        headers: [String: String] = [:]
+    ) async throws -> (Data, HTTPURLResponse) {
         let url = URL(string: "http://\(host):\(port)/invoke")!
         var request = URLRequest(url: url)
         request.httpMethod = "POST"
         request.setValue("application/json", forHTTPHeaderField: "Content-Type")
+        for (key, value) in headers {
+            request.setValue(value, forHTTPHeaderField: key)
+        }
         request.httpBody = payload.data(using: .utf8)
         request.timeoutInterval = 10.0