Lambda-Runtime-Invocation-Id

.autover/changes/4133a18f-b6f9-4f24-886e-217756fbb670.json

@@ -0,0 +1,11 @@
+{
+  "Projects": [
+    {
+      "Name": "Amazon.Lambda.RuntimeSupport",
+      "Type": "Minor",
+      "ChangelogMessages": [
+        "Add support for the Lambda-Runtime-Invocation-Id header for cross-wiring invoke protection. The runtime echoes the header back on the response and error calls when the Runtime API provides it, and treats an HTTP 410 Gone (invoke timeout) response as a non-fatal condition, logging it and continuing to the next invocation."
+      ]
+    }
+  ]
+}

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs

@@ -380,11 +380,17 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
 
             Func<Task> processingFunc = async () =>
             {
+                // Per-invocation id echoed back on /response and /error for cross-wiring protection.
+                // Null when the Runtime API did not send the Lambda-Runtime-Invocation-Id header, in
+                // which case nothing is echoed. Captured as a local so concurrent invocations in
+                // multi-concurrency mode each carry their own id.
+                string invocationId = null;
                 if (invocation.LambdaContext is LambdaContext impl)
                 {
                     Client.ConsoleLogger.SetRuntimeHeaders(impl.RuntimeApiHeaders);
                     SetInvocationTraceId(impl.RuntimeApiHeaders.TraceId);
                     SetSerializerOnContext(impl);
+                    invocationId = impl.InvocationId;
                 }
 
                 // Initialize ResponseStreamFactory — includes RuntimeApiClient reference
@@ -395,7 +401,8 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
                         invocation.LambdaContext.AwsRequestId,
                         isMultiConcurrency,
                         runtimeApiClient,
-                        cancellationToken);
+                        cancellationToken,
+                        invocationId);
                 }
 
                 try
@@ -420,7 +427,7 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
                         }
                         else
                         {
-                        await Client.ReportInvocationErrorAsync(invocation.LambdaContext.AwsRequestId, exception, cancellationToken);
+                        await Client.ReportInvocationErrorAsync(invocation.LambdaContext.AwsRequestId, invocationId, exception, cancellationToken);
                     }
                     }
                     finally
@@ -450,7 +457,7 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
                         _logger.LogInformation("Starting sending response");
                         try
                         {
-                            await Client.SendResponseAsync(invocation.LambdaContext.AwsRequestId, response?.OutputStream, cancellationToken);
+                            await Client.SendResponseAsync(invocation.LambdaContext.AwsRequestId, invocationId, response?.OutputStream, cancellationToken);
                         }
                         finally
                         {
@@ -465,6 +472,15 @@ internal async Task InvokeOnceAsync(CancellationToken cancellationToken = defaul
 
                     _logger.LogInformation("Finished InvokeOnceAsync");
                 }
+                catch (RuntimeApiInvokeTimeoutException timeout)
+                {
+                    // The Runtime API rejected the response or error with HTTP 410 Gone because the
+                    // invocation had already timed out (cross-wiring protection). This is expected and
+                    // not fatal — the response would have been discarded anyway. Log via the internal
+                    // logger (not the customer log) and loop back to /next in both on-demand and
+                    // multi-concurrency modes rather than crashing the process.
+                    _logger.LogInformation($"Invocation {timeout.AwsRequestId} timed out before its response was submitted; the Runtime API rejected it. Continuing to the next invocation.");
+                }
                 catch(Exception ex)
                 {
                     // Only capture and continue for multi concurrency because we do not want to change

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/RawStreamingHttpClient.cs

@@ -62,11 +62,13 @@ public RawStreamingHttpClient(string hostAndPort)
         /// for error reporting.
         /// </summary>
         /// <param name="awsRequestId">The Lambda request ID.</param>
+        /// <param name="invocationId">The unique-per-invocation id to echo back for cross-wiring protection. When null, the header is not sent.</param>
         /// <param name="responseStream">The response stream that provides data and error state.</param>
         /// <param name="userAgent">The User-Agent header value.</param>
         /// <param name="cancellationToken">Cancellation token.</param>
         public async Task SendStreamingResponseAsync(
             string awsRequestId,
+            string invocationId,
             ResponseStream responseStream,
             string userAgent,
             CancellationToken cancellationToken = default)
@@ -86,6 +88,10 @@ public async Task SendStreamingResponseAsync(
             headers.Append($"{StreamingConstants.ResponseModeHeader}: {StreamingConstants.StreamingResponseMode}\r\n");
             headers.Append("Transfer-Encoding: chunked\r\n");
             headers.Append($"Trailer: {StreamingConstants.ErrorTypeTrailer}, {StreamingConstants.ErrorBodyTrailer}\r\n");
+            // Echo the per-invocation id back for cross-wiring protection. Only sent when the
+            // Runtime API provided it on /next.
+            if (!string.IsNullOrEmpty(invocationId))
+                headers.Append($"{RuntimeApiHeaders.HeaderInvocationId}: {invocationId}\r\n");
             headers.Append("\r\n");
 
             var headerBytes = Encoding.ASCII.GetBytes(headers.ToString());

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/ResponseStreamContext.cs

@@ -30,6 +30,12 @@ internal class ResponseStreamContext
         /// </summary>
         public string AwsRequestId { get; set; }
 
+        /// <summary>
+        /// The unique-per-invocation id echoed back on the streaming response for cross-wiring
+        /// protection. Null when the Runtime API did not provide it, in which case nothing is echoed.
+        /// </summary>
+        public string InvocationId { get; set; }
+
         /// <summary>
         /// Whether CreateStream() has been called for this invocation.
         /// </summary>

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/ResponseStreamFactory.cs

@@ -62,7 +62,7 @@ public static ResponseStream CreateStream(byte[] prelude)
             // This runs concurrently — SerializeToStreamAsync will block
             // until the handler finishes writing or reports an error.
             context.SendTask = context.RuntimeApiClient.StartStreamingResponseAsync(
-                context.AwsRequestId, lambdaStream, context.CancellationToken);
+                context.AwsRequestId, context.InvocationId, lambdaStream, context.CancellationToken);
 
             return lambdaStream;
         }
@@ -71,11 +71,13 @@ public static ResponseStream CreateStream(byte[] prelude)
 
         internal static void InitializeInvocation(
             string awsRequestId, bool isMultiConcurrency,
-            RuntimeApiClient runtimeApiClient, CancellationToken cancellationToken)
+            RuntimeApiClient runtimeApiClient, CancellationToken cancellationToken,
+            string invocationId = null)
         {
             var context = new ResponseStreamContext
             {
                 AwsRequestId = awsRequestId,
+                InvocationId = invocationId,
                 StreamCreated = false,
                 Stream = null,
                 RuntimeApiClient = runtimeApiClient,

Libraries/src/Amazon.Lambda.RuntimeSupport/Client/IRuntimeApiClient.cs

@@ -65,7 +65,19 @@ public interface IRuntimeApiClient
         /// <param name="cancellationToken">The optional cancellation token to use.</param>
         /// <returns>A Task representing the asynchronous operation.</returns>
         Task ReportInvocationErrorAsync(string awsRequestId, Exception exception, CancellationToken cancellationToken = default);
-
+
+        /// <summary>
+        /// Report an invocation error as an asynchronous operation, echoing the invocation id for
+        /// cross-wiring protection.
+        /// </summary>
+        /// <param name="awsRequestId">The ID of the function request that caused the error.</param>
+        /// <param name="invocationId">The unique-per-invocation id to echo back to the Runtime API. When null, the header is not sent.</param>
+        /// <param name="exception">The exception to report.</param>
+        /// <param name="cancellationToken">The optional cancellation token to use.</param>
+        /// <returns>A Task representing the asynchronous operation.</returns>
+        /// <exception cref="RuntimeApiInvokeTimeoutException">The invocation timed out before the error was submitted (cross-wiring protection).</exception>
+        Task ReportInvocationErrorAsync(string awsRequestId, string invocationId, Exception exception, CancellationToken cancellationToken = default);
+
         /// <summary>
         ///  Triggers the snapshot to be taken, and then after resume, restores the lambda
         /// context from the Runtime API as an asynchronous operation when SnapStart is enabled.
@@ -91,5 +103,17 @@ public interface IRuntimeApiClient
         /// <param name="cancellationToken">The optional cancellation token to use.</param>
         /// <returns></returns>
         Task SendResponseAsync(string awsRequestId, Stream outputStream, CancellationToken cancellationToken = default);
+
+        /// <summary>
+        /// Send a response to a function invocation to the Runtime API as an asynchronous operation,
+        /// echoing the invocation id for cross-wiring protection.
+        /// </summary>
+        /// <param name="awsRequestId">The ID of the function request being responded to.</param>
+        /// <param name="invocationId">The unique-per-invocation id to echo back to the Runtime API. When null, the header is not sent.</param>
+        /// <param name="outputStream">The content of the response to the function invocation.</param>
+        /// <param name="cancellationToken">The optional cancellation token to use.</param>
+        /// <returns></returns>
+        /// <exception cref="RuntimeApiInvokeTimeoutException">The invocation timed out before the response was submitted (cross-wiring protection).</exception>
+        Task SendResponseAsync(string awsRequestId, string invocationId, Stream outputStream, CancellationToken cancellationToken = default);
     }
 }