Merge branch 'normj/response-streaming' into normj/aspnetcore-responsestreaming

Author: normj

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

@@ -14,6 +14,7 @@
  */
 
 using System;
+using System.Buffers;
 using System.IO;
 using System.Text;
 using System.Threading;
@@ -94,10 +95,25 @@ private async Task WritePreludeAsync(CancellationToken cancellationToken = defau
                         ThrowIfCompletedOrError();
                     }
 
-                    await _httpOutputStream.WriteAsync(_prelude, 0, _prelude.Length, cancellationToken);
-                    await _httpOutputStream.WriteAsync(PreludeDelimiter, 0, PreludeDelimiter.Length, cancellationToken);
+                    // Combine prelude and delimiter into a single write so they arrive
+                    // as one contiguous chunk. Writing them separately would produce two
+                    // chunked-encoding frames, and the Runtime API / API Gateway expects
+                    // the prelude JSON immediately followed by the 8-byte null delimiter
+                    // as a single unit to parse correctly.
+                    var combinedLength = _prelude.Length + PreludeDelimiter.Length;
+                    var combined = ArrayPool<byte>.Shared.Rent(combinedLength);
+                    try
+                    {
+                        Buffer.BlockCopy(_prelude, 0, combined, 0, _prelude.Length);
+                        Buffer.BlockCopy(PreludeDelimiter, 0, combined, _prelude.Length, PreludeDelimiter.Length);
 
-                    await _httpOutputStream.FlushAsync(cancellationToken);
+                        await _httpOutputStream.WriteAsync(combined, 0, combinedLength, cancellationToken);
+                        await _httpOutputStream.FlushAsync(cancellationToken);
+                    }
+                    finally
+                    {
+                        ArrayPool<byte>.Shared.Return(combined);
+                    }
                 }
                 finally
                 {

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/ResponseStreamTests.cs

@@ -284,6 +284,107 @@ public async Task SetHttpOutputStreamAsync_NullPrelude_WritesNoPreludeBytes()
             Assert.Empty(output.ToArray());
         }
 
+        // ---- Prelude + delimiter single-chunk tests (via ChunkedStreamWriter) ----
+
+        [Fact]
+        public async Task SetHttpOutputStreamAsync_WithPrelude_ViaChunkedWriter_ProducesSingleChunk()
+        {
+            var preludeJson = Encoding.UTF8.GetBytes("{\"statusCode\":200}");
+            var rs = new ResponseStream(preludeJson);
+            var rawOutput = new MemoryStream();
+            var chunkedWriter = new ChunkedStreamWriter(rawOutput);
+
+            await rs.SetHttpOutputStreamAsync(chunkedWriter);
+
+            var wireBytes = Encoding.ASCII.GetString(rawOutput.ToArray());
+
+            // The prelude (18 bytes) + delimiter (8 bytes) = 26 bytes = 0x1A
+            // Should be exactly one chunk: "1A\r\n{prelude}{8 null bytes}\r\n"
+            var expectedDataLength = preludeJson.Length + 8; // 26
+            var expectedHex = expectedDataLength.ToString("X");
+            Assert.StartsWith($"{expectedHex}\r\n", wireBytes);
+
+            // Verify there is only one chunk header (only one hex size prefix)
+            var chunkCount = 0;
+            var remaining = wireBytes;
+            while (remaining.Length > 0)
+            {
+                var crlfIndex = remaining.IndexOf("\r\n", StringComparison.Ordinal);
+                if (crlfIndex < 0) break;
+                var sizeStr = remaining.Substring(0, crlfIndex);
+                if (int.TryParse(sizeStr, System.Globalization.NumberStyles.HexNumber, null, out var chunkSize) && chunkSize >= 0)
+                {
+                    chunkCount++;
+                    // Skip past: hex\r\n{data}\r\n
+                    remaining = remaining.Substring(crlfIndex + 2 + chunkSize + 2);
+                }
+                else
+                {
+                    break;
+                }
+            }
+            Assert.Equal(1, chunkCount);
+        }
+
+        [Fact]
+        public async Task SetHttpOutputStreamAsync_WithPrelude_ViaChunkedWriter_DelimiterImmediatelyFollowsPrelude()
+        {
+            var preludeJson = Encoding.UTF8.GetBytes("{\"statusCode\":201}");
+            var rs = new ResponseStream(preludeJson);
+            var rawOutput = new MemoryStream();
+            var chunkedWriter = new ChunkedStreamWriter(rawOutput);
+
+            await rs.SetHttpOutputStreamAsync(chunkedWriter);
+
+            // Parse the chunk to get the raw data payload
+            var wireBytes = rawOutput.ToArray();
+            var wireStr = Encoding.ASCII.GetString(wireBytes);
+            var firstCrlf = wireStr.IndexOf("\r\n", StringComparison.Ordinal);
+            var dataStart = firstCrlf + 2;
+            var dataLength = preludeJson.Length + 8;
+            var chunkData = new byte[dataLength];
+            Array.Copy(wireBytes, dataStart, chunkData, 0, dataLength);
+
+            // First part should be the prelude JSON
+            Assert.Equal(preludeJson, chunkData[..preludeJson.Length]);
+            // Immediately followed by 8 null bytes (delimiter)
+            Assert.Equal(new byte[8], chunkData[preludeJson.Length..]);
+        }
+
+        [Fact]
+        public async Task SetHttpOutputStreamAsync_WithPrelude_ViaChunkedWriter_HandlerDataInSeparateChunk()
+        {
+            var preludeJson = Encoding.UTF8.GetBytes("{\"statusCode\":200}");
+            var rs = new ResponseStream(preludeJson);
+            var rawOutput = new MemoryStream();
+            var chunkedWriter = new ChunkedStreamWriter(rawOutput);
+
+            await rs.SetHttpOutputStreamAsync(chunkedWriter);
+            await rs.WriteAsync(Encoding.UTF8.GetBytes("body data"), 0, 9);
+
+            var wireStr = Encoding.ASCII.GetString(rawOutput.ToArray());
+
+            // Should have exactly 2 chunks: one for prelude+delimiter, one for body
+            var chunkCount = 0;
+            var remaining = wireStr;
+            while (remaining.Length > 0)
+            {
+                var crlfIndex = remaining.IndexOf("\r\n", StringComparison.Ordinal);
+                if (crlfIndex < 0) break;
+                var sizeStr = remaining.Substring(0, crlfIndex);
+                if (int.TryParse(sizeStr, System.Globalization.NumberStyles.HexNumber, null, out var chunkSize) && chunkSize >= 0)
+                {
+                    chunkCount++;
+                    remaining = remaining.Substring(crlfIndex + 2 + chunkSize + 2);
+                }
+                else
+                {
+                    break;
+                }
+            }
+            Assert.Equal(2, chunkCount);
+        }
+
         // ---- MarkCompleted idempotency ----
 
         [Fact]