Correctly propagate the trace ID

Package.swift

@@ -36,6 +36,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", from: "1.10.1"),
         .package(url: "https://github.com/apple/swift-collections.git", from: "1.3.0"),
         .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.10.1"),
+        .package(url: "https://github.com/apple/swift-service-context.git", from: "1.3.0"),
     ],
     targets: [
         .target(
@@ -44,6 +45,7 @@ let package = Package(
                 .product(name: "NIOCore", package: "swift-nio"),
                 .product(name: "DequeModule", package: "swift-collections"),
                 .product(name: "Logging", package: "swift-log"),
+                .product(name: "ServiceContextModule", package: "swift-service-context"),
                 .product(name: "NIOHTTP1", package: "swift-nio"),
                 .product(name: "NIOPosix", package: "swift-nio"),
                 .product(
@@ -74,6 +76,7 @@ let package = Package(
             name: "AWSLambdaRuntimeTests",
             dependencies: [
                 .byName(name: "AWSLambdaRuntime"),
+                .product(name: "ServiceContextModule", package: "swift-service-context"),
                 .product(name: "NIOTestUtils", package: "swift-nio"),
                 .product(name: "NIOFoundationCompat", package: "swift-nio"),
             ],

Package@swift-6.0.swift

@@ -25,6 +25,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", from: "1.8.0"),
         .package(url: "https://github.com/apple/swift-collections.git", from: "1.3.0"),
         .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.9.0"),
+        .package(url: "https://github.com/apple/swift-service-context.git", from: "1.3.0"),
     ],
     targets: [
         .target(
@@ -35,6 +36,7 @@ let package = Package(
                 .product(name: "Logging", package: "swift-log"),
                 .product(name: "NIOHTTP1", package: "swift-nio"),
                 .product(name: "NIOPosix", package: "swift-nio"),
+                .product(name: "ServiceContextModule", package: "swift-service-context"),
                 .product(name: "ServiceLifecycle", package: "swift-service-lifecycle"),
             ],
             swiftSettings: defaultSwiftSettings
@@ -61,6 +63,7 @@ let package = Package(
                 .byName(name: "AWSLambdaRuntime"),
                 .product(name: "NIOTestUtils", package: "swift-nio"),
                 .product(name: "NIOFoundationCompat", package: "swift-nio"),
+                .product(name: "ServiceContextModule", package: "swift-service-context"),
             ],
             swiftSettings: defaultSwiftSettings
         ),

Sources/AWSLambdaRuntime/Lambda.swift

@@ -17,6 +17,7 @@ import Dispatch
 import Logging
 import NIOCore
 import NIOPosix
+import ServiceContextModule
 
 #if os(macOS)
 import Darwin.C
@@ -90,12 +91,13 @@ public enum Lambda {
 
                 logger.trace("Waiting for next invocation")
                 let (invocation, writer) = try await runtimeClient.nextInvocation()
+                let traceId = invocation.metadata.traceID
 
                 // Create a per-request logger with request-specific metadata
                 let requestLogger = loggingConfiguration.makeLogger(
                     label: "Lambda",
                     requestID: invocation.metadata.requestID,
-                    traceID: invocation.metadata.traceID
+                    traceID: traceId
                 )
 
                 // when log level is trace or lower, print the first 6 Mb of the payload
@@ -116,26 +118,46 @@ public enum Lambda {
                     metadata: metadata
                 )
 
-                do {
-                    try await handler.handle(
-                        invocation.event,
-                        responseWriter: writer,
-                        context: LambdaContext(
-                            requestID: invocation.metadata.requestID,
-                            traceID: invocation.metadata.traceID,
-                            tenantID: invocation.metadata.tenantID,
-                            invokedFunctionARN: invocation.metadata.invokedFunctionARN,
-                            deadline: LambdaClock.Instant(
-                                millisecondsSinceEpoch: invocation.metadata.deadlineInMillisSinceEpoch
-                            ),
-                            logger: requestLogger
+                // Wrap handler invocation in a ServiceContext scope so that
+                // downstream libraries can access the trace ID via
+                // ServiceContext.current?.traceID without depending on AWSLambdaRuntime.
+                // In single-concurrency mode, also set the _X_AMZN_TRACE_ID env var
+                // for backward compatibility with legacy tooling.
+                var serviceContext = ServiceContext.current ?? ServiceContext.topLevel
+                serviceContext.traceID = traceId
+                try await ServiceContext.withValue(serviceContext) {
+                    if isSingleConcurrencyMode {
+                        setenv("_X_AMZN_TRACE_ID", traceId, 1)
+                    }
+                    defer {
+                        if isSingleConcurrencyMode {
+                            unsetenv("_X_AMZN_TRACE_ID")
+                        }
+                    }
+
+                    do {
+                        try await handler.handle(
+                            invocation.event,
+                            responseWriter: writer,
+                            context: LambdaContext(
+                                requestID: invocation.metadata.requestID,
+                                traceID: traceId,
+                                tenantID: invocation.metadata.tenantID,
+                                invokedFunctionARN: invocation.metadata.invokedFunctionARN,
+                                deadline: LambdaClock.Instant(
+                                    millisecondsSinceEpoch: invocation.metadata.deadlineInMillisSinceEpoch
+                                ),
+                                logger: requestLogger
+                            )
+                        )
+                        requestLogger.trace("Handler finished processing invocation")
+                    } catch {
+                        requestLogger.trace(
+                            "Handler failed processing invocation",
+                            metadata: ["Handler error": "\(error)"]
                         )
-                    )
-                    requestLogger.trace("Handler finished processing invocation")
-                } catch {
-                    requestLogger.trace("Handler failed processing invocation", metadata: ["Handler error": "\(error)"])
-                    try await writer.reportError(error)
-                    continue
+                        try await writer.reportError(error)
+                    }
                 }
             }
         } catch is CancellationError {

Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift

@@ -0,0 +1,50 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 ServiceContextModule
+
+// MARK: - ServiceContext integration
+
+/// A ``ServiceContextKey`` for the AWS X-Ray trace ID.
+///
+/// This allows downstream libraries that depend on `swift-service-context`
+/// (but not on `AWSLambdaRuntime`) to access the current trace ID via
+/// `ServiceContext.current?.traceID`.
+private enum LambdaTraceIDKey: ServiceContextKey {
+    typealias Value = String
+    static var nameOverride: String? { AmazonHeaders.traceID }
+}
+
+extension ServiceContext {
+    /// The AWS X-Ray trace ID for the current Lambda invocation, if available.
+    ///
+    /// This value is automatically set by the Lambda runtime before calling the handler
+    /// and is available to all code running within the handler's async task tree.
+    ///
+    /// Downstream libraries can read this without depending on `AWSLambdaRuntime`:
+    /// ```swift
+    /// if let traceID = ServiceContext.current?.traceID {
+    ///     // propagate traceID to outgoing HTTP requests, etc.
+    /// }
+    /// ```
+    public var traceID: String? {
+        get {
+            self[LambdaTraceIDKey.self]
+        }
+        set {
+            self[LambdaTraceIDKey.self] = newValue
+        }
+    }
+}