From 3b99bef34d8572541383f20fb20bf2b35d76b875 Mon Sep 17 00:00:00 2001
From: Jonas Ha <61319894+j-d-ha@users.noreply.github.com>
Date: Wed, 10 Dec 2025 09:55:22 -0500
Subject: [PATCH 01/29] fix(ci): remove sign-commits from changelog workflow
 (#223)

* chore(changelog): update for v1.3.0 release

- Add new features, bug fixes, and refactoring notes.
- Document maintenance updates and improvements.
- Prepare for the next release iteration.

* chore(workflows): update changelog workflow settings

- Remove `sign-commits` option from changelog workflow.
- Adjust minor formatting inconsistency in `id` field.
---
 .github/workflows/update-changelog.yaml |  3 +--
 CHANGELOG.md                            | 26 ++++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/update-changelog.yaml b/.github/workflows/update-changelog.yaml
index a6b7d41d..4b85125e 100644
--- a/.github/workflows/update-changelog.yaml
+++ b/.github/workflows/update-changelog.yaml
@@ -33,7 +33,7 @@ jobs:
         run: perl -i -pe 'BEGIN{undef $/} s/(## (?:\[)?v[\d.]+(?:\](?:\([^)]+\))?)? - \d{4}-\d{2}-\d{2}\n)[\s\S]*?### Changes\n+/$1\n/g' CHANGELOG.md
 
       - name: Create Pull Request
-        id: cpr
+        id: cpr 
         uses: peter-evans/create-pull-request@v7
         with:
           branch: changelog/${{ github.event.release.tag_name }}
@@ -48,7 +48,6 @@ jobs:
           commit-message: "chore: update changelog for ${{ github.event.release.tag_name }}"
           labels: "automated,changelog,chore"
           delete-branch: true
-          sign-commits: true
 
       - name: Enable Auto-merge
         if: steps.cpr.outputs.pull-request-number
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 76c82c0a..54899f36 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,31 @@
 All notable changes to this project will be documented in this file.
 See [Conventional Commits](https://conventionalcommits.org) for commit guidelines.
 
-## [Unreleased](https://github.com/j-d-ha/aws-lambda-host/compare/v1.2.1...HEAD)
+## [Unreleased](https://github.com/j-d-ha/aws-lambda-host/compare/v1.3.0...HEAD)
+
+## [v1.3.0](https://github.com/j-d-ha/aws-lambda-host/compare/v1.2.1...v1.3.0) - 2025-12-10
+
+## 🚀 Features
+
+* feat(source-generators): support multiple MapHandler invocations with custom feature providers (#214) @j-d-ha
+* docs: update MkDocs palette toggle configuration (#211) @j-d-ha
+
+## 🐛 Bug Fixes
+
+* fix: update third-party license attributions (#217) @j-d-ha
+
+## 📚 Documentation
+
+* docs: add comprehensive getting started guide and restructure documentation (#209) @j-d-ha
+
+## 🔄 Refactoring
+
+* refactor(host): migrate BootstrapHttpClient from options to dependency injection (#219) @j-d-ha
+* refactor(docs): replace ASPNETCORE_ENVIRONMENT with DOTNET_ENVIRONMENT (#216) @j-d-ha
+
+## 🔧 Maintenance
+
+* ci(github): optimize workflow triggers for draft PRs (#215) @j-d-ha
 
 ## [v1.2.1](https://github.com/j-d-ha/aws-lambda-host/compare/v1.2.0...v1.2.1) - 2025-11-30
 

From e990b93a0dce72b77b477608c8dd475eb61adcf7 Mon Sep 17 00:00:00 2001
From: Jonas Ha <61319894+j-d-ha@users.noreply.github.com>
Date: Wed, 10 Dec 2025 16:15:46 -0500
Subject: [PATCH 02/29] feat(context): add TenantId and TraceId properties to
 DefaultLambdaHostContext (#224)

- Expose TenantId and TraceId from the lambda context.
- Helps in identifying tenant-specific and traceable context data.
---
 src/AwsLambda.Host/Core/Context/DefaultLambdaHostContext.cs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/AwsLambda.Host/Core/Context/DefaultLambdaHostContext.cs b/src/AwsLambda.Host/Core/Context/DefaultLambdaHostContext.cs
index b405ebe9..4c07e181 100644
--- a/src/AwsLambda.Host/Core/Context/DefaultLambdaHostContext.cs
+++ b/src/AwsLambda.Host/Core/Context/DefaultLambdaHostContext.cs
@@ -69,6 +69,10 @@ public async ValueTask DisposeAsync()
 
     public TimeSpan RemainingTime => _lambdaContext.RemainingTime;
 
+    public string TenantId => _lambdaContext.TenantId;
+
+    public string TraceId => _lambdaContext.TraceId;
+
     //      ┌──────────────────────────────────────────────────────────┐
     //      │                    ILambdaHostContext                    │
     //      └──────────────────────────────────────────────────────────┘

From 305f8b28452957ede0905d6a30d09856eb2370ef Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Wed, 10 Dec 2025 16:29:56 -0500
Subject: [PATCH 03/29] chore: update changelog for v1.3.1 (#225)

Co-authored-by: j-d-ha <61319894+j-d-ha@users.noreply.github.com>
---
 CHANGELOG.md | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 54899f36..2ffe12b9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,14 @@
 All notable changes to this project will be documented in this file.
 See [Conventional Commits](https://conventionalcommits.org) for commit guidelines.
 
-## [Unreleased](https://github.com/j-d-ha/aws-lambda-host/compare/v1.3.0...HEAD)
+## [Unreleased](https://github.com/j-d-ha/aws-lambda-host/compare/v1.3.1...HEAD)
+
+## [v1.3.1](https://github.com/j-d-ha/aws-lambda-host/compare/v1.3.0...v1.3.1) - 2025-12-10
+
+### 🐛 Bug Fixes
+
+* fix(core): add missing TenantId and TraceId properties to DefaultLambdaHostContext (#224) @j-d-ha
+* fix(ci): remove sign-commits from changelog workflow (#223) @j-d-ha
 
 ## [v1.3.0](https://github.com/j-d-ha/aws-lambda-host/compare/v1.2.1...v1.3.0) - 2025-12-10
 

From 4adfcff3984e3f9d50d39eca4ce2c1fb73ceacf0 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Sun, 7 Dec 2025 17:16:56 -0500
Subject: [PATCH 04/29] refactor(testing): replace LambdaClient with
 InvocationProcessor and enhance test server

- Replaced `LambdaClient` with `InvocationProcessor` for improved invocation handling and clarity.
- Updated `LambdaTestServer` to manage lifecycle and invocation processing via `InvocationProcessor`.
- Removed `LambdaClient` creation logic from `LambdaApplicationFactory`.
- Consolidated redundant methods in `LambdaTestServer` for a streamlined design.
- Introduced the `ServerState` enum to represent `LambdaTestServer`'s state transitions.
---
 .../Tests/LambdaHostTest.cs                   |  55 +-
 .../Internal/InvocationProcessor.cs           | 403 +++++++++++++
 .../LambdaApplicationFactory.cs               |  51 +-
 .../LambdaTestClient.cs                       | 142 -----
 .../LambdaTestServer.cs                       | 539 +++++++-----------
 src/AwsLambda.Host.Testing/ServerState.cs     |  37 ++
 6 files changed, 681 insertions(+), 546 deletions(-)
 create mode 100644 src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
 delete mode 100644 src/AwsLambda.Host.Testing/LambdaTestClient.cs
 create mode 100644 src/AwsLambda.Host.Testing/ServerState.cs

diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index bcab7c64..26a01fb1 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -13,10 +13,9 @@ public class LambdaHostTest
     public async Task LambdaHost_CanStartWithoutError()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
-        var client = factory.CreateClient();
-        // No need to wait for next request - server handles this automatically
-        var response = await client.InvokeAsync<string, string>(
+        var response = await factory.Server.InvokeAsync<string, string>(
             "Jonas",
             TestContext.Current.CancellationToken
         );
@@ -41,9 +40,9 @@ public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
             );
         });
 
-        var client = factory.CreateClient();
-        // No need to wait for next request - server handles this automatically
-        var response = await client.InvokeAsync<string, string>(
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        var response = await factory.Server.InvokeAsync<string, string>(
             "Jonas",
             TestContext.Current.CancellationToken
         );
@@ -56,13 +55,13 @@ public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
     public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
-        var client = factory.CreateClient();
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
         // Launch 5 concurrent invocations
         var tasks = Enumerable
             .Range(1, 5)
             .Select(i =>
-                client.InvokeAsync<string, string>(
+                factory.Server.InvokeAsync<string, string>(
                     $"User{i}",
                     TestContext.Current.CancellationToken
                 )
@@ -84,9 +83,9 @@ public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
     public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
-        var client = factory.CreateClient();
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
-        var response = await client.InvokeAsync<string, int>(
+        var response = await factory.Server.InvokeAsync<string, int>(
             123,
             TestContext.Current.CancellationToken
         );
@@ -100,38 +99,30 @@ public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
     public async Task InvokeAsync_WithPreCanceledToken_CancelsInvocation()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
-        var client = factory.CreateClient();
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
         using var cts = new CancellationTokenSource();
         await cts.CancelAsync();
 
         await Assert.ThrowsAsync<TaskCanceledException>(() =>
-            client.InvokeAsync<string, string>("Jonas", cts.Token)
+            factory.Server.InvokeAsync<string, string>("Jonas", cts.Token)
         );
     }
 
     [Fact]
     public async Task InvokeAsync_WithZeroTimeout_CancelsInvocation() =>
-        await Assert.ThrowsAsync<AggregateException>(async () =>
+        await Assert.ThrowsAsync<TaskCanceledException>(async () =>
         {
-            try
-            {
-                await using var factory = new LambdaApplicationFactory<Program>();
-                var client = factory
-                    .CreateClient()
-                    .ConfigureOptions(options =>
-                        options.InvocationHeaderOptions.ClientWaitTimeout = TimeSpan.Zero
-                    );
-
-                await client.InvokeAsync<string, string>(
-                    "Jonas",
-                    TestContext.Current.CancellationToken
-                );
-            }
-            catch (Exception e)
-            {
-                Console.WriteLine(e.GetType().FullName);
-                throw;
-            }
+            await using var factory = new LambdaApplicationFactory<Program>();
+            await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+            var options = new LambdaClientOptions();
+            options.InvocationHeaderOptions.ClientWaitTimeout = TimeSpan.Zero;
+
+            await factory.Server.InvokeAsync<string, string>(
+                "Jonas",
+                options,
+                TestContext.Current.CancellationToken
+            );
         });
 }
diff --git a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
new file mode 100644
index 00000000..3f596427
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
@@ -0,0 +1,403 @@
+using System.Collections.Concurrent;
+using System.Net;
+using System.Text;
+using System.Text.Json;
+using System.Threading.Channels;
+using Microsoft.AspNetCore.Routing;
+
+namespace AwsLambda.Host.Testing;
+
+/// <summary>
+/// Internal processor that handles HTTP transactions from Lambda Bootstrap.
+/// Routes requests, queues invocations, and manages request-response correlation.
+/// </summary>
+internal class InvocationProcessor : IAsyncDisposable
+{
+    private const int TransactionChannelCapacity = 1024;
+    private const int NextRequestChannelCapacity = 1024;
+
+    private static readonly OperationCanceledException DisposedException = new(
+        "InvocationProcessor disposed"
+    );
+
+    private readonly JsonSerializerOptions _jsonSerializerOptions;
+    private readonly ConcurrentQueue<string> _pendingInvocationIds;
+    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
+    private readonly Channel<LambdaHttpTransaction> _queuedNextRequests;
+    private readonly ILambdaRuntimeRouteManager _routeManager;
+    private readonly CancellationTokenSource _shutdownCts;
+    private readonly Channel<LambdaHttpTransaction> _transactionChannel;
+    private Task? _processingTask;
+
+    internal InvocationProcessor(
+        JsonSerializerOptions? jsonSerializerOptions = null,
+        ILambdaRuntimeRouteManager? routeManager = null
+    )
+    {
+        _transactionChannel = Channel.CreateBounded<LambdaHttpTransaction>(
+            new BoundedChannelOptions(TransactionChannelCapacity)
+            {
+                SingleReader = true,
+                SingleWriter = false,
+                FullMode = BoundedChannelFullMode.Wait,
+            }
+        );
+        _pendingInvocationIds = new ConcurrentQueue<string>();
+        _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
+        _queuedNextRequests = Channel.CreateBounded<LambdaHttpTransaction>(
+            new BoundedChannelOptions(NextRequestChannelCapacity)
+            {
+                SingleReader = false,
+                SingleWriter = false,
+                FullMode = BoundedChannelFullMode.Wait,
+            }
+        );
+        _routeManager = routeManager ?? new LambdaRuntimeRouteManager();
+        _jsonSerializerOptions = jsonSerializerOptions ?? new JsonSerializerOptions();
+        _shutdownCts = new CancellationTokenSource();
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _shutdownCts.CancelAsync();
+
+        _transactionChannel.Writer.TryComplete();
+        _queuedNextRequests.Writer.TryComplete();
+
+        // Cancel any transactions waiting for work
+        while (_queuedNextRequests.Reader.TryRead(out var queuedTransaction))
+            queuedTransaction.Fail(DisposedException);
+
+        // Fail any in-flight transactions that haven't been processed yet
+        while (_transactionChannel.Reader.TryRead(out var transaction))
+            transaction.Fail(DisposedException);
+
+        // Cancel any pending invocations waiting for bootstrap responses
+        foreach (var pendingInvocation in _pendingInvocations.Values)
+            pendingInvocation.ResponseTcs.TrySetCanceled(_shutdownCts.Token);
+
+        if (_processingTask != null)
+            try
+            {
+                await _processingTask.ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                // Expected when task is canceled
+            }
+
+        _shutdownCts.Dispose();
+    }
+
+    /// <summary>
+    /// Creates the HTTP handler for Lambda Bootstrap to use.
+    /// </summary>
+    internal HttpMessageHandler CreateTestingHandler() =>
+        new LambdaTestingHttpHandler(_transactionChannel);
+
+    /// <summary>
+    /// Starts the background processing loop.
+    /// Called by LambdaTestServer after host starts.
+    /// </summary>
+    internal void StartProcessing()
+    {
+        if (_processingTask != null)
+            throw new InvalidOperationException("Processor already started");
+
+        _processingTask = Task.Run(ProcessTransactionsAsync);
+    }
+
+    /// <summary>
+    /// Queues a new invocation to be processed by Lambda Bootstrap.
+    /// Called by LambdaTestServer.InvokeAsync() or LambdaClient.InvokeAsync().
+    /// </summary>
+    internal async Task<InvocationCompletion> QueueInvocationAsync(
+        string requestId,
+        HttpResponseMessage eventResponse,
+        DateTimeOffset deadlineUtc,
+        CancellationToken cancellationToken
+    )
+    {
+        var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
+
+        if (!_pendingInvocations.TryAdd(requestId, pending))
+            throw new InvalidOperationException($"Duplicate request ID: {requestId}");
+
+        _pendingInvocationIds.Enqueue(requestId);
+
+        // If there's a queued /next request, serve it immediately
+        if (_queuedNextRequests.Reader.TryRead(out var nextTransaction))
+            RespondToNextRequest(nextTransaction);
+
+        using var cancellationRegistration = cancellationToken.Register(() =>
+            CancelPendingInvocation(requestId, cancellationToken)
+        );
+
+        // Wait for Bootstrap to process and respond
+        return await pending.ResponseTcs.Task.WaitAsync(cancellationToken);
+    }
+
+    /// <summary>
+    /// Background loop that processes transactions from the handler.
+    /// </summary>
+    private async Task ProcessTransactionsAsync()
+    {
+        await foreach (
+            var transaction in _transactionChannel.Reader.ReadAllAsync(_shutdownCts.Token)
+        )
+            try
+            {
+                await ProcessTransactionAsync(transaction);
+            }
+            catch (Exception ex)
+            {
+                // Fail the transaction and continue processing
+                transaction.Fail(ex);
+            }
+    }
+
+    /// <summary>
+    /// Routes a single transaction based on request type.
+    /// </summary>
+    private async Task ProcessTransactionAsync(LambdaHttpTransaction transaction)
+    {
+        if (!_routeManager.TryMatch(transaction.Request, out var requestType, out var routeValues))
+            throw new InvalidOperationException(
+                $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
+            );
+
+        switch (requestType!.Value)
+        {
+            case RequestType.GetNextInvocation:
+                await HandleGetNextInvocationAsync(transaction);
+                break;
+
+            case RequestType.PostResponse:
+                await HandlePostResponseAsync(transaction, routeValues!);
+                break;
+
+            case RequestType.PostError:
+                await HandlePostErrorAsync(transaction, routeValues!);
+                break;
+
+            default:
+                throw new InvalidOperationException(
+                    $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
+                );
+        }
+    }
+
+    /// <summary>
+    /// Handles GET /invocation/next - Bootstrap polling for work.
+    /// </summary>
+    private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)
+    {
+        // Try to dequeue next pending invocation (FIFO)
+        if (!TryDequeuePendingInvocation(out var pending))
+        {
+            // No work available - queue this /next request
+            try
+            {
+                await _queuedNextRequests.Writer.WriteAsync(transaction, _shutdownCts.Token);
+            }
+            catch (OperationCanceledException)
+            {
+                transaction.Fail(DisposedException);
+            }
+
+            return;
+        }
+
+        RespondToNextRequest(transaction, pending);
+    }
+
+    /// <summary>
+    /// Responds to a /next request with a pending invocation.
+    /// </summary>
+    private void RespondToNextRequest(LambdaHttpTransaction transaction, PendingInvocation pending)
+    {
+        // Respond with the event payload and Lambda headers
+        if (transaction.Respond(pending.EventResponse))
+            return;
+
+        // Request was already canceled; re-enqueue invocation to avoid dropping it
+        _pendingInvocationIds.Enqueue(pending.RequestId);
+    }
+
+    /// <summary>
+    /// Responds to a /next request with a pending invocation by looking up the next one.
+    /// </summary>
+    private void RespondToNextRequest(LambdaHttpTransaction transaction)
+    {
+        // Try to dequeue next pending invocation (FIFO)
+        if (!TryDequeuePendingInvocation(out var pending))
+        {
+            // This shouldn't happen, but if it does, respond with error
+            transaction.Respond(
+                new HttpResponseMessage(HttpStatusCode.InternalServerError)
+                {
+                    Content = new StringContent("No pending invocations available"),
+                }
+            );
+            return;
+        }
+
+        RespondToNextRequest(transaction, pending);
+    }
+
+    /// <summary>
+    /// Handles POST /invocation/{requestId}/response - successful function execution.
+    /// </summary>
+    private async Task HandlePostResponseAsync(
+        LambdaHttpTransaction transaction,
+        RouteValueDictionary routeValues
+    )
+    {
+        var requestId = routeValues["requestId"]?.ToString();
+
+        if (
+            string.IsNullOrEmpty(requestId)
+            || !_pendingInvocations.TryRemove(requestId, out var pending)
+        )
+        {
+            transaction.Respond(
+                new HttpResponseMessage(HttpStatusCode.NotFound)
+                {
+                    Content = new StringContent(
+                        string.IsNullOrEmpty(requestId)
+                            ? "Missing requestId"
+                            : $"No pending invocation for request ID: {requestId}"
+                    ),
+                }
+            );
+            return;
+        }
+
+        // Complete the invocation with the response from Bootstrap
+        pending.ResponseTcs.SetResult(
+            await CreateCompletionAsync(RequestType.PostResponse, transaction.Request)
+        );
+
+        // Acknowledge to Bootstrap
+        transaction.Respond(
+            new HttpResponseMessage(HttpStatusCode.Accepted)
+            {
+                Content = new StringContent(
+                    JsonSerializer.Serialize(
+                        new Dictionary<string, string> { ["status"] = "success" },
+                        _jsonSerializerOptions
+                    ),
+                    Encoding.UTF8,
+                    "application/json"
+                ),
+                Version = Version.Parse("1.1"),
+            }
+        );
+    }
+
+    /// <summary>
+    /// Handles POST /invocation/{requestId}/error - function execution failed.
+    /// </summary>
+    private async Task HandlePostErrorAsync(
+        LambdaHttpTransaction transaction,
+        RouteValueDictionary routeValues
+    )
+    {
+        var requestId = routeValues["requestId"]?.ToString();
+
+        if (
+            string.IsNullOrEmpty(requestId)
+            || !_pendingInvocations.TryRemove(requestId, out var pending)
+        )
+        {
+            transaction.Respond(
+                new HttpResponseMessage(HttpStatusCode.NotFound)
+                {
+                    Content = new StringContent(
+                        string.IsNullOrEmpty(requestId)
+                            ? "Missing requestId"
+                            : $"No pending invocation for request ID: {requestId}"
+                    ),
+                }
+            );
+            return;
+        }
+
+        // Complete the invocation with the error response from Bootstrap
+        pending.ResponseTcs.SetResult(
+            await CreateCompletionAsync(RequestType.PostError, transaction.Request)
+        );
+
+        // Acknowledge to Bootstrap
+        transaction.Respond(
+            new HttpResponseMessage(HttpStatusCode.Accepted)
+            {
+                Content = new StringContent(
+                    "{\"status\":\"success\"}",
+                    Encoding.UTF8,
+                    "application/json"
+                ),
+                Version = Version.Parse("1.1"),
+            }
+        );
+    }
+
+    private bool TryDequeuePendingInvocation(out PendingInvocation pendingInvocation)
+    {
+        var now = DateTimeOffset.UtcNow;
+
+        while (_pendingInvocationIds.TryDequeue(out var requestId))
+            if (_pendingInvocations.TryGetValue(requestId, out pendingInvocation))
+            {
+                if (pendingInvocation.DeadlineUtc <= now)
+                {
+                    if (_pendingInvocations.TryRemove(requestId, out var expiredInvocation))
+                        expiredInvocation.ResponseTcs.TrySetCanceled();
+
+                    continue;
+                }
+
+                return true;
+            }
+
+        pendingInvocation = null!;
+        return false;
+    }
+
+    private void CancelPendingInvocation(string requestId, CancellationToken cancellationToken)
+    {
+        if (_pendingInvocations.TryRemove(requestId, out var pendingInvocation))
+            pendingInvocation.ResponseTcs.TrySetCanceled(cancellationToken);
+    }
+
+    private static async Task<InvocationCompletion> CreateCompletionAsync(
+        RequestType requestType,
+        HttpRequestMessage sourceRequest
+    )
+    {
+        var clonedRequest = new HttpRequestMessage(sourceRequest.Method, sourceRequest.RequestUri)
+        {
+            Version = sourceRequest.Version,
+            VersionPolicy = sourceRequest.VersionPolicy,
+        };
+
+        foreach (var header in sourceRequest.Headers)
+            clonedRequest.Headers.TryAddWithoutValidation(header.Key, header.Value);
+
+        foreach (var option in sourceRequest.Options)
+            clonedRequest.Options.TryAdd(option.Key, option.Value);
+
+        if (sourceRequest.Content != null)
+        {
+            var contentBytes = await sourceRequest.Content.ReadAsByteArrayAsync();
+            var clonedContent = new ByteArrayContent(contentBytes);
+
+            foreach (var header in sourceRequest.Content.Headers)
+                clonedContent.Headers.TryAddWithoutValidation(header.Key, header.Value);
+
+            clonedRequest.Content = clonedContent;
+        }
+
+        return new InvocationCompletion { Request = clonedRequest, RequestType = requestType };
+    }
+}
diff --git a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
index 440371dc..609eefe8 100644
--- a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
+++ b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
@@ -73,7 +73,7 @@ public class LambdaApplicationFactory<TEntryPoint> : IDisposable, IAsyncDisposab
     /// <summary>
     /// Gets the <see cref="LambdaTestServer"/> created by this <see cref="LambdaApplicationFactory{TEntryPoint}"/>.
     /// </summary>
-    internal LambdaTestServer Server
+    public LambdaTestServer Server
     {
         get
         {
@@ -85,14 +85,7 @@ internal LambdaTestServer Server
     /// <summary>
     /// Gets the <see cref="IServiceProvider"/> created by the server associated with this <see cref="LambdaApplicationFactory{TEntryPoint}"/>.
     /// </summary>
-    public virtual IServiceProvider Services
-    {
-        get
-        {
-            EnsureServer();
-            return _host!.Services;
-        }
-    }
+    public virtual IServiceProvider Services => Server.Services;
 
     /// <inheritdoc />
     public virtual async ValueTask DisposeAsync()
@@ -106,15 +99,10 @@ public virtual async ValueTask DisposeAsync()
         foreach (var factory in _derivedFactories)
             await ((IAsyncDisposable)factory).DisposeAsync().ConfigureAwait(false);
 
+        // Server handles disposing both processor and host
         if (_server != null)
             await _server.DisposeAsync().ConfigureAwait(false);
 
-        if (_host != null)
-        {
-            await _host.StopAsync().ConfigureAwait(false);
-            _host?.Dispose();
-        }
-
         _disposedAsync = true;
 
         Dispose(true);
@@ -134,12 +122,6 @@ public void Dispose()
     /// </summary>
     ~LambdaApplicationFactory() => Dispose(false);
 
-    public LambdaClient CreateClient()
-    {
-        EnsureServer();
-        return _server!.CreateLambdaClient();
-    }
-
     /// <summary>
     /// Creates a new <see cref="LambdaApplicationFactory{TEntryPoint}"/> with a <see cref="IHostBuilder"/>
     /// that is further customized by <paramref name="configuration"/>.
@@ -152,15 +134,12 @@ public LambdaApplicationFactory<TEntryPoint> WithHostBuilder(
         Action<IHostBuilder> configuration
     ) => WithHostBuilderCore(configuration);
 
-    internal virtual LambdaTestServer CreateServer() => new();
-
     internal virtual LambdaApplicationFactory<TEntryPoint> WithHostBuilderCore(
         Action<IHostBuilder> configuration
     )
     {
         var factory = new DelegatedLambdaApplicationFactory(
             ClientOptions,
-            CreateServer,
             CreateHost,
             CreateHostBuilder,
             GetTestAssemblies,
@@ -231,12 +210,12 @@ private void ConfigureHostBuilder(IHostBuilder hostBuilder)
         SetContentRoot(hostBuilder);
         _configuration(hostBuilder);
 
-        _server = CreateServer();
+        var processor = new InvocationProcessor();
 
         // set Lambda Bootstrap Http Client
         hostBuilder.ConfigureServices(services =>
         {
-            services.AddLambdaBootstrapHttpClient(new HttpClient(_server.CreateTestingHandler()));
+            services.AddLambdaBootstrapHttpClient(new HttpClient(processor.CreateTestingHandler()));
 
             services.PostConfigure<LambdaHostOptions>(options =>
             {
@@ -245,10 +224,11 @@ private void ConfigureHostBuilder(IHostBuilder hostBuilder)
             });
         });
 
+        // Build the host but DON'T start it - server will start it
         _host = CreateHost(hostBuilder);
 
-        // Start the server's background processing loop after host is ready
-        _server.Start();
+        // Create the public server with the built (but not started) host
+        _server = new LambdaTestServer(_host, processor);
     }
 
     private void SetContentRoot(IHostBuilder builder)
@@ -440,12 +420,9 @@ private static void EnsureDepsFile()
     /// </summary>
     /// <param name="builder">The <see cref="IHostBuilder"/> used to create the host.</param>
     /// <returns>The <see cref="IHost"/> with the bootstrapped application.</returns>
-    protected virtual IHost CreateHost(IHostBuilder builder)
-    {
-        var host = builder.Build();
-        host.Start();
-        return host;
-    }
+    protected virtual IHost CreateHost(IHostBuilder builder) =>
+        // Build the host but DON'T start it - LambdaTestServer.StartAsync() will start it
+        builder.Build();
 
     /// <summary>
     /// Gives a fixture an opportunity to configure the application before it gets built.
@@ -475,12 +452,10 @@ private sealed class DelegatedLambdaApplicationFactory : LambdaApplicationFactor
     {
         private readonly Func<IHostBuilder, IHost> _createHost;
         private readonly Func<IHostBuilder?> _createHostBuilder;
-        private readonly Func<LambdaTestServer> _createServer;
         private readonly Func<IEnumerable<Assembly>> _getTestAssemblies;
 
         public DelegatedLambdaApplicationFactory(
             LambdaApplicationFactoryClientOptions options,
-            Func<LambdaTestServer> createServer,
             Func<IHostBuilder, IHost> createHost,
             Func<IHostBuilder?> createHostBuilder,
             Func<IEnumerable<Assembly>> getTestAssemblies,
@@ -488,7 +463,6 @@ Action<IHostBuilder> configureWebHost
         )
         {
             ClientOptions = options;
-            _createServer = createServer;
             _createHost = createHost;
             _createHostBuilder = createHostBuilder;
             _getTestAssemblies = getTestAssemblies;
@@ -497,8 +471,6 @@ Action<IHostBuilder> configureWebHost
 
         protected override IHost CreateHost(IHostBuilder builder) => _createHost(builder);
 
-        internal override LambdaTestServer CreateServer() => _createServer();
-
         protected override IHostBuilder? CreateHostBuilder() => _createHostBuilder();
 
         protected override IEnumerable<Assembly> GetTestAssemblies() => _getTestAssemblies();
@@ -510,7 +482,6 @@ Action<IHostBuilder> configuration
         ) =>
             new DelegatedLambdaApplicationFactory(
                 ClientOptions,
-                _createServer,
                 _createHost,
                 _createHostBuilder,
                 _getTestAssemblies,
diff --git a/src/AwsLambda.Host.Testing/LambdaTestClient.cs b/src/AwsLambda.Host.Testing/LambdaTestClient.cs
deleted file mode 100644
index cb370a24..00000000
--- a/src/AwsLambda.Host.Testing/LambdaTestClient.cs
+++ /dev/null
@@ -1,142 +0,0 @@
-using System.Net;
-using System.Net.Http.Json;
-using System.Text;
-using System.Text.Json;
-
-namespace AwsLambda.Host.Testing;
-
-/// <summary>
-/// Client for invoking Lambda functions in tests.
-/// Provides a clean API that abstracts HTTP details.
-/// </summary>
-public class LambdaClient
-{
-    private readonly JsonSerializerOptions _jsonSerializerOptions;
-    private readonly LambdaClientOptions _lambdaClientOptions;
-    private readonly LambdaTestServer _server;
-    private int _requestCounter;
-
-    internal LambdaClient(LambdaTestServer server, JsonSerializerOptions jsonSerializerOptions)
-    {
-        _server = server;
-        _jsonSerializerOptions = jsonSerializerOptions;
-        _lambdaClientOptions = new LambdaClientOptions();
-    }
-
-    /// <summary>
-    /// Configures client options for invocation headers.
-    /// </summary>
-    public LambdaClient ConfigureOptions(Action<LambdaClientOptions> configureOptions)
-    {
-        ArgumentNullException.ThrowIfNull(configureOptions);
-
-        configureOptions(_lambdaClientOptions);
-
-        return this;
-    }
-
-    /// <summary>
-    /// Invokes the Lambda function with the given event and waits for the response.
-    /// </summary>
-    public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
-        TEvent invokeEvent,
-        CancellationToken cancellationToken = default
-    )
-    {
-        // Generate unique request ID
-        var requestId = GetRequestId();
-
-        // Create the event response with Lambda headers
-        var eventResponse = CreateEventResponse(invokeEvent, requestId);
-        var deadlineUtc = DateTimeOffset.UtcNow.Add(
-            _lambdaClientOptions.InvocationHeaderOptions.ClientWaitTimeout
-        );
-
-        // Queue invocation and wait for Bootstrap to process it
-        var waitTimeout = _lambdaClientOptions.InvocationHeaderOptions.ClientWaitTimeout;
-
-        using var timeoutCts = new CancellationTokenSource(waitTimeout);
-        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(
-            cancellationToken,
-            timeoutCts.Token
-        );
-
-        var completion = await _server.QueueInvocationAsync(
-            requestId,
-            eventResponse,
-            deadlineUtc,
-            linkedCts.Token
-        );
-
-        var responseMessage = completion.Request;
-        var wasSuccess = completion.RequestType == RequestType.PostResponse;
-
-        var invocationResponse = new InvocationResponse<TResponse>
-        {
-            WasSuccess = wasSuccess,
-            Response = wasSuccess
-                ? await (
-                    responseMessage.Content?.ReadFromJsonAsync<TResponse>(
-                        _jsonSerializerOptions,
-                        cancellationToken
-                    ) ?? Task.FromResult<TResponse?>(default)
-                )
-                : default,
-            Error = !wasSuccess
-                ? await (
-                    responseMessage.Content?.ReadFromJsonAsync<ErrorResponse>(
-                        _jsonSerializerOptions,
-                        cancellationToken
-                    ) ?? Task.FromResult<ErrorResponse?>(null)
-                )
-                : null,
-        };
-
-        return invocationResponse;
-    }
-
-    private HttpResponseMessage CreateEventResponse<TEvent>(TEvent invokeEvent, string requestId)
-    {
-        var response = new HttpResponseMessage(HttpStatusCode.OK)
-        {
-            Content = new StringContent(
-                JsonSerializer.Serialize(invokeEvent, _jsonSerializerOptions),
-                Encoding.UTF8,
-                "application/json"
-            ),
-            Version = Version.Parse("1.1"),
-        };
-
-        // Add standard HTTP headers
-        response.Headers.Date = new DateTimeOffset(
-            _lambdaClientOptions.InvocationHeaderOptions.Date
-        );
-        response.Headers.TransferEncodingChunked = _lambdaClientOptions
-            .InvocationHeaderOptions
-            .TransferEncodingChunked;
-
-        // Add custom Lambda runtime headers
-        var deadlineMs = DateTimeOffset
-            .UtcNow.Add(_lambdaClientOptions.InvocationHeaderOptions.FunctionTimeout)
-            .ToUnixTimeMilliseconds();
-        response.Headers.Add("Lambda-Runtime-Deadline-Ms", deadlineMs.ToString());
-        response.Headers.Add("Lambda-Runtime-Aws-Request-Id", requestId);
-        response.Headers.Add(
-            "Lambda-Runtime-Trace-Id",
-            _lambdaClientOptions.InvocationHeaderOptions.TraceId
-        );
-        response.Headers.Add(
-            "Lambda-Runtime-Invoked-Function-Arn",
-            _lambdaClientOptions.InvocationHeaderOptions.FunctionArn
-        );
-
-        // Add any additional custom headers
-        foreach (var header in _lambdaClientOptions.InvocationHeaderOptions.AdditionalHeaders)
-            response.Headers.Add(header.Key, header.Value);
-
-        return response;
-    }
-
-    private string GetRequestId() =>
-        Interlocked.Increment(ref _requestCounter).ToString().PadLeft(12, '0');
-}
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServer.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
index 773dceb3..7e11ac3e 100644
--- a/src/AwsLambda.Host.Testing/LambdaTestServer.cs
+++ b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
@@ -1,419 +1,294 @@
-using System.Collections.Concurrent;
 using System.Net;
+using System.Net.Http.Json;
 using System.Text;
 using System.Text.Json;
-using System.Threading.Channels;
-using Microsoft.AspNetCore.Routing;
+using Microsoft.Extensions.Hosting;
 
 namespace AwsLambda.Host.Testing;
 
 /// <summary>
-/// Test server that processes HTTP transactions from Lambda Bootstrap.
-/// Routes requests, queues invocations, and manages request-response correlation.
+/// Test server that manages the Lambda host lifecycle and invocation processing.
+/// Provides explicit StartAsync/StopAsync control and direct invocation capabilities.
 /// </summary>
-internal class LambdaTestServer : IAsyncDisposable
+public class LambdaTestServer : IAsyncDisposable
 {
-    private const int TransactionChannelCapacity = 1024;
-    private const int NextRequestChannelCapacity = 1024;
-    private static readonly OperationCanceledException DisposedException = new(
-        "LambdaTestServer disposed"
-    );
-
-    private readonly LambdaClient _client;
+    private readonly IHost _host;
+    private readonly InvocationProcessor _processor;
+    private readonly LambdaClientOptions _defaultClientOptions;
     private readonly JsonSerializerOptions _jsonSerializerOptions;
-    private readonly ConcurrentQueue<string> _pendingInvocationIds;
-    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
-    private readonly Channel<LambdaHttpTransaction> _queuedNextRequests;
-    private readonly ILambdaRuntimeRouteManager _routeManager;
-    private readonly CancellationTokenSource _shutdownCts;
-    private readonly Channel<LambdaHttpTransaction> _transactionChannel;
-    private Task? _processingTask;
+    private int _requestCounter;
+    private ServerState _state;
 
     internal LambdaTestServer(
-        JsonSerializerOptions? jsonSerializerOptions = null,
-        ILambdaRuntimeRouteManager? routeManager = null
+        IHost host,
+        InvocationProcessor processor,
+        LambdaClientOptions? defaultClientOptions = null,
+        JsonSerializerOptions? jsonSerializerOptions = null
     )
     {
-        _transactionChannel = Channel.CreateBounded<LambdaHttpTransaction>(
-            new BoundedChannelOptions(TransactionChannelCapacity)
-            {
-                SingleReader = true,
-                SingleWriter = false,
-                FullMode = BoundedChannelFullMode.Wait,
-            }
-        );
-        _pendingInvocationIds = new ConcurrentQueue<string>();
-        _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
-        _queuedNextRequests = Channel.CreateBounded<LambdaHttpTransaction>(
-            new BoundedChannelOptions(NextRequestChannelCapacity)
-            {
-                SingleReader = false,
-                SingleWriter = false,
-                FullMode = BoundedChannelFullMode.Wait,
-            }
-        );
-        _routeManager = routeManager ?? new LambdaRuntimeRouteManager();
+        _host = host ?? throw new ArgumentNullException(nameof(host));
+        _processor = processor ?? throw new ArgumentNullException(nameof(processor));
+        _defaultClientOptions = defaultClientOptions ?? new LambdaClientOptions();
         _jsonSerializerOptions = jsonSerializerOptions ?? new JsonSerializerOptions();
-        _shutdownCts = new CancellationTokenSource();
-
-        // Create client that communicates with this server
-        _client = new LambdaClient(this, _jsonSerializerOptions);
-    }
-
-    public async ValueTask DisposeAsync()
-    {
-        await _shutdownCts.CancelAsync();
-
-        _transactionChannel.Writer.TryComplete();
-        _queuedNextRequests.Writer.TryComplete();
-
-        // Cancel any transactions waiting for work
-        while (_queuedNextRequests.Reader.TryRead(out var queuedTransaction))
-            queuedTransaction.Fail(DisposedException);
-
-        // Fail any in-flight transactions that haven't been processed yet
-        while (_transactionChannel.Reader.TryRead(out var transaction))
-            transaction.Fail(DisposedException);
-
-        // Cancel any pending invocations waiting for bootstrap responses
-        foreach (var pendingInvocation in _pendingInvocations.Values)
-            pendingInvocation.ResponseTcs.TrySetCanceled(_shutdownCts.Token);
-
-        if (_processingTask != null)
-        {
-            try
-            {
-                await _processingTask.ConfigureAwait(false);
-            }
-            catch (OperationCanceledException)
-            {
-                // Expected when task is canceled
-            }
-        }
-
-        _shutdownCts.Dispose();
+        _state = ServerState.Created;
     }
 
     /// <summary>
-    /// Creates the HTTP handler for Lambda Bootstrap to use.
+    /// Gets the current lifecycle state of the server.
     /// </summary>
-    internal HttpMessageHandler CreateTestingHandler() =>
-        new LambdaTestingHttpHandler(_transactionChannel);
+    public ServerState State => _state;
 
     /// <summary>
-    /// Gets the client for test code to invoke Lambda functions.
+    /// Gets the service provider from the host.
+    /// Only available after the server has been started.
     /// </summary>
-    internal LambdaClient CreateLambdaClient() => _client;
-
-    /// <summary>
-    /// Starts the background processing loop.
-    /// Called automatically by LambdaApplicationFactory after host starts.
-    /// </summary>
-    internal void Start()
+    /// <exception cref="InvalidOperationException">Thrown if server is not in Running state.</exception>
+    public IServiceProvider Services
     {
-        if (_processingTask != null)
-            throw new InvalidOperationException("Server already started");
+        get
+        {
+            if (_state != ServerState.Running)
+                throw new InvalidOperationException(
+                    $"Services are only available when server is Running. Current state: {_state}. Call StartAsync() first."
+                );
 
-        _processingTask = Task.Run(ProcessTransactionsAsync);
+            return _host.Services;
+        }
     }
 
     /// <summary>
-    /// Queues a new invocation to be processed by Lambda Bootstrap.
-    /// Called by LambdaClient.InvokeAsync().
+    /// Starts the Lambda host and begins processing invocations.
     /// </summary>
-    internal async Task<InvocationCompletion> QueueInvocationAsync(
-        string requestId,
-        HttpResponseMessage eventResponse,
-        DateTimeOffset deadlineUtc,
-        CancellationToken cancellationToken
-    )
+    /// <param name="cancellationToken">Cancellation token.</param>
+    /// <exception cref="InvalidOperationException">Thrown if server is not in Created state.</exception>
+    public async Task StartAsync(CancellationToken cancellationToken = default)
     {
-        var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
-
-        if (!_pendingInvocations.TryAdd(requestId, pending))
-            throw new InvalidOperationException($"Duplicate request ID: {requestId}");
+        if (_state != ServerState.Created)
+            throw new InvalidOperationException(
+                $"Server can only be started from Created state. Current state: {_state}"
+            );
 
-        _pendingInvocationIds.Enqueue(requestId);
+        try
+        {
+            _state = ServerState.Starting;
 
-        // If there's a queued /next request, serve it immediately
-        if (_queuedNextRequests.Reader.TryRead(out var nextTransaction))
-            RespondToNextRequest(nextTransaction);
+            // Start the host
+            await _host.StartAsync(cancellationToken);
 
-        using var cancellationRegistration = cancellationToken.Register(() =>
-            CancelPendingInvocation(requestId, cancellationToken)
-        );
+            // Start background processing
+            _processor.StartProcessing();
 
-        // Wait for Bootstrap to process and respond
-        return await pending.ResponseTcs.Task.WaitAsync(cancellationToken);
+            _state = ServerState.Running;
+        }
+        catch
+        {
+            // Cleanup on failure
+            await DisposeAsync();
+            throw;
+        }
     }
 
     /// <summary>
-    /// Background loop that processes transactions from the handler.
+    /// Stops the Lambda host.
     /// </summary>
-    private async Task ProcessTransactionsAsync()
+    /// <param name="cancellationToken">Cancellation token.</param>
+    public async Task StopAsync(CancellationToken cancellationToken = default)
     {
-        await foreach (
-            var transaction in _transactionChannel.Reader.ReadAllAsync(_shutdownCts.Token)
-        )
-            try
-            {
-                await ProcessTransactionAsync(transaction);
-            }
-            catch (Exception ex)
-            {
-                // Fail the transaction and continue processing
-                transaction.Fail(ex);
-            }
-    }
+        // Idempotent - allow multiple calls
+        if (_state == ServerState.Stopped || _state == ServerState.Disposed)
+            return;
 
-    /// <summary>
-    /// Routes a single transaction based on request type.
-    /// </summary>
-    private async Task ProcessTransactionAsync(LambdaHttpTransaction transaction)
-    {
-        if (!_routeManager.TryMatch(transaction.Request, out var requestType, out var routeValues))
+        if (_state != ServerState.Running)
             throw new InvalidOperationException(
-                $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
+                $"Server can only be stopped from Running state. Current state: {_state}"
             );
 
-        switch (requestType!.Value)
+        try
         {
-            case RequestType.GetNextInvocation:
-                await HandleGetNextInvocationAsync(transaction);
-                break;
-
-            case RequestType.PostResponse:
-                await HandlePostResponseAsync(transaction, routeValues!);
-                break;
-
-            case RequestType.PostError:
-                await HandlePostErrorAsync(transaction, routeValues!);
-                break;
-
-            default:
-                throw new InvalidOperationException(
-                    $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
-                );
+            _state = ServerState.Stopping;
+            await _host.StopAsync(cancellationToken);
+            _state = ServerState.Stopped;
         }
-    }
-
-    /// <summary>
-    /// Handles GET /invocation/next - Bootstrap polling for work.
-    /// </summary>
-    private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)
-    {
-        // Try to dequeue next pending invocation (FIFO)
-        if (!TryDequeuePendingInvocation(out var pending))
+        catch (OperationCanceledException)
         {
-            // No work available - queue this /next request
-            try
-            {
-                await _queuedNextRequests.Writer.WriteAsync(transaction, _shutdownCts.Token);
-            }
-            catch (OperationCanceledException)
-            {
-                transaction.Fail(DisposedException);
-            }
-
-            return;
+            // Still mark as stopped even if shutdown was canceled
+            _state = ServerState.Stopped;
+            throw;
         }
-
-        RespondToNextRequest(transaction, pending);
     }
 
     /// <summary>
-    /// Responds to a /next request with a pending invocation.
+    /// Invokes the Lambda function with the given event and waits for the response.
     /// </summary>
-    private void RespondToNextRequest(LambdaHttpTransaction transaction, PendingInvocation pending)
-    {
-        // Respond with the event payload and Lambda headers
-        if (transaction.Respond(pending.EventResponse))
-            return;
-
-        // Request was already canceled; re-enqueue invocation to avoid dropping it
-        _pendingInvocationIds.Enqueue(pending.RequestId);
-    }
+    /// <typeparam name="TResponse">The type of the response.</typeparam>
+    /// <typeparam name="TEvent">The type of the event.</typeparam>
+    /// <param name="invokeEvent">The event to invoke with.</param>
+    /// <param name="cancellationToken">Cancellation token.</param>
+    /// <returns>The invocation response.</returns>
+    /// <exception cref="InvalidOperationException">Thrown if server is not in Running state.</exception>
+    public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
+        TEvent invokeEvent,
+        CancellationToken cancellationToken = default
+    ) => await InvokeAsync<TResponse, TEvent>(invokeEvent, null, cancellationToken);
 
     /// <summary>
-    /// Responds to a /next request with a pending invocation by looking up the next one.
+    /// Invokes the Lambda function with the given event and waits for the response.
     /// </summary>
-    private void RespondToNextRequest(LambdaHttpTransaction transaction)
+    /// <typeparam name="TResponse">The type of the response.</typeparam>
+    /// <typeparam name="TEvent">The type of the event.</typeparam>
+    /// <param name="invokeEvent">The event to invoke with.</param>
+    /// <param name="options">Client options for this invocation.</param>
+    /// <param name="cancellationToken">Cancellation token.</param>
+    /// <returns>The invocation response.</returns>
+    /// <exception cref="InvalidOperationException">Thrown if server is not in Running state.</exception>
+    public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
+        TEvent invokeEvent,
+        LambdaClientOptions? options,
+        CancellationToken cancellationToken = default
+    )
     {
-        // Try to dequeue next pending invocation (FIFO)
-        if (!TryDequeuePendingInvocation(out var pending))
-        {
-            // This shouldn't happen, but if it does, respond with error
-            transaction.Respond(
-                new HttpResponseMessage(HttpStatusCode.InternalServerError)
-                {
-                    Content = new StringContent("No pending invocations available"),
-                }
+        if (_state != ServerState.Running)
+            throw new InvalidOperationException(
+                _state == ServerState.Created
+                    ? "Server must be started before invoking. Call StartAsync() first."
+                    : $"Server is {_state}. Only Running servers can process invocations."
             );
-            return;
-        }
 
-        RespondToNextRequest(transaction, pending);
-    }
+        // Use provided options or fall back to defaults
+        var clientOptions = options ?? _defaultClientOptions;
 
-    /// <summary>
-    /// Handles POST /invocation/{requestId}/response - successful function execution.
-    /// </summary>
-    private async Task HandlePostResponseAsync(
-        LambdaHttpTransaction transaction,
-        RouteValueDictionary routeValues
-    )
-    {
-        var requestId = routeValues["requestId"]?.ToString();
+        // Generate unique request ID
+        var requestId = GetRequestId();
 
-        if (
-            string.IsNullOrEmpty(requestId)
-            || !_pendingInvocations.TryRemove(requestId, out var pending)
-        )
-        {
-            transaction.Respond(
-                new HttpResponseMessage(HttpStatusCode.NotFound)
-                {
-                    Content = new StringContent(
-                        string.IsNullOrEmpty(requestId)
-                            ? "Missing requestId"
-                            : $"No pending invocation for request ID: {requestId}"
-                    ),
-                }
-            );
-            return;
-        }
+        // Create the event response with Lambda headers
+        var eventResponse = CreateEventResponse(invokeEvent, requestId, clientOptions);
+        var deadlineUtc = DateTimeOffset.UtcNow.Add(
+            clientOptions.InvocationHeaderOptions.ClientWaitTimeout
+        );
 
-        // Complete the invocation with the response from Bootstrap
-        pending.ResponseTcs.SetResult(
-            await CreateCompletionAsync(RequestType.PostResponse, transaction.Request)
+        // Queue invocation and wait for Bootstrap to process it
+        var waitTimeout = clientOptions.InvocationHeaderOptions.ClientWaitTimeout;
+
+        using var timeoutCts = new CancellationTokenSource(waitTimeout);
+        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(
+            cancellationToken,
+            timeoutCts.Token
         );
 
-        // Acknowledge to Bootstrap
-        transaction.Respond(
-            new HttpResponseMessage(HttpStatusCode.Accepted)
-            {
-                Content = new StringContent(
-                    JsonSerializer.Serialize(
-                        new Dictionary<string, string> { ["status"] = "success" },
-                        _jsonSerializerOptions
-                    ),
-                    Encoding.UTF8,
-                    "application/json"
-                ),
-                Version = Version.Parse("1.1"),
-            }
+        var completion = await _processor.QueueInvocationAsync(
+            requestId,
+            eventResponse,
+            deadlineUtc,
+            linkedCts.Token
         );
 
-        return;
+        var responseMessage = completion.Request;
+        var wasSuccess = completion.RequestType == RequestType.PostResponse;
+
+        var invocationResponse = new InvocationResponse<TResponse>
+        {
+            WasSuccess = wasSuccess,
+            Response = wasSuccess
+                ? await (
+                    responseMessage.Content?.ReadFromJsonAsync<TResponse>(
+                        _jsonSerializerOptions,
+                        cancellationToken
+                    ) ?? Task.FromResult<TResponse?>(default)
+                )
+                : default,
+            Error = !wasSuccess
+                ? await (
+                    responseMessage.Content?.ReadFromJsonAsync<ErrorResponse>(
+                        _jsonSerializerOptions,
+                        cancellationToken
+                    ) ?? Task.FromResult<ErrorResponse?>(null)
+                )
+                : null,
+        };
+
+        return invocationResponse;
     }
 
     /// <summary>
-    /// Handles POST /invocation/{requestId}/error - function execution failed.
+    /// Disposes the server, stopping the host if needed.
     /// </summary>
-    private async Task HandlePostErrorAsync(
-        LambdaHttpTransaction transaction,
-        RouteValueDictionary routeValues
-    )
+    public async ValueTask DisposeAsync()
     {
-        var requestId = routeValues["requestId"]?.ToString();
+        if (_state == ServerState.Disposed)
+            return;
 
-        if (
-            string.IsNullOrEmpty(requestId)
-            || !_pendingInvocations.TryRemove(requestId, out var pending)
-        )
+        try
         {
-            transaction.Respond(
-                new HttpResponseMessage(HttpStatusCode.NotFound)
+            // Stop if running
+            if (_state == ServerState.Running)
+                try
                 {
-                    Content = new StringContent(
-                        string.IsNullOrEmpty(requestId)
-                            ? "Missing requestId"
-                            : $"No pending invocation for request ID: {requestId}"
-                    ),
+                    await StopAsync();
                 }
-            );
-            return;
-        }
-
-        // Complete the invocation with the error response from Bootstrap
-        pending.ResponseTcs.SetResult(
-            await CreateCompletionAsync(RequestType.PostError, transaction.Request)
-        );
-
-        // Acknowledge to Bootstrap
-        transaction.Respond(
-            new HttpResponseMessage(HttpStatusCode.Accepted)
-            {
-                Content = new StringContent(
-                    "{\"status\":\"success\"}",
-                    Encoding.UTF8,
-                    "application/json"
-                ),
-                Version = Version.Parse("1.1"),
-            }
-        );
-
-        return;
-    }
-
-    private bool TryDequeuePendingInvocation(out PendingInvocation pendingInvocation)
-    {
-        var now = DateTimeOffset.UtcNow;
-
-        while (_pendingInvocationIds.TryDequeue(out var requestId))
-        {
-            if (_pendingInvocations.TryGetValue(requestId, out pendingInvocation))
-            {
-                if (pendingInvocation.DeadlineUtc <= now)
+                catch
                 {
-                    if (_pendingInvocations.TryRemove(requestId, out var expiredInvocation))
-                        expiredInvocation.ResponseTcs.TrySetCanceled();
-
-                    continue;
+                    // Best effort - continue with disposal
                 }
 
-                return true;
-            }
-        }
-
-        pendingInvocation = null!;
-        return false;
-    }
+            // Dispose processor
+            await _processor.DisposeAsync();
 
-    private void CancelPendingInvocation(string requestId, CancellationToken cancellationToken)
-    {
-        if (_pendingInvocations.TryRemove(requestId, out var pendingInvocation))
-            pendingInvocation.ResponseTcs.TrySetCanceled(cancellationToken);
+            // Dispose host
+            if (_host is IAsyncDisposable asyncDisposable)
+                await asyncDisposable.DisposeAsync();
+            else
+                _host.Dispose();
+        }
+        finally
+        {
+            _state = ServerState.Disposed;
+        }
     }
 
-    private static async Task<InvocationCompletion> CreateCompletionAsync(
-        RequestType requestType,
-        HttpRequestMessage sourceRequest
+    private HttpResponseMessage CreateEventResponse<TEvent>(
+        TEvent invokeEvent,
+        string requestId,
+        LambdaClientOptions clientOptions
     )
     {
-        var clonedRequest = new HttpRequestMessage(sourceRequest.Method, sourceRequest.RequestUri)
+        var response = new HttpResponseMessage(HttpStatusCode.OK)
         {
-            Version = sourceRequest.Version,
-            VersionPolicy = sourceRequest.VersionPolicy,
+            Content = new StringContent(
+                JsonSerializer.Serialize(invokeEvent, _jsonSerializerOptions),
+                Encoding.UTF8,
+                "application/json"
+            ),
+            Version = Version.Parse("1.1"),
         };
 
-        foreach (var header in sourceRequest.Headers)
-            clonedRequest.Headers.TryAddWithoutValidation(header.Key, header.Value);
-
-        foreach (var option in sourceRequest.Options)
-            clonedRequest.Options.TryAdd(option.Key, option.Value);
-
-        if (sourceRequest.Content != null)
-        {
-            var contentBytes = await sourceRequest.Content.ReadAsByteArrayAsync();
-            var clonedContent = new ByteArrayContent(contentBytes);
-
-            foreach (var header in sourceRequest.Content.Headers)
-                clonedContent.Headers.TryAddWithoutValidation(header.Key, header.Value);
+        // Add standard HTTP headers
+        response.Headers.Date = new DateTimeOffset(clientOptions.InvocationHeaderOptions.Date);
+        response.Headers.TransferEncodingChunked = clientOptions
+            .InvocationHeaderOptions
+            .TransferEncodingChunked;
+
+        // Add custom Lambda runtime headers
+        var deadlineMs = DateTimeOffset
+            .UtcNow.Add(clientOptions.InvocationHeaderOptions.FunctionTimeout)
+            .ToUnixTimeMilliseconds();
+        response.Headers.Add("Lambda-Runtime-Deadline-Ms", deadlineMs.ToString());
+        response.Headers.Add("Lambda-Runtime-Aws-Request-Id", requestId);
+        response.Headers.Add(
+            "Lambda-Runtime-Trace-Id",
+            clientOptions.InvocationHeaderOptions.TraceId
+        );
+        response.Headers.Add(
+            "Lambda-Runtime-Invoked-Function-Arn",
+            clientOptions.InvocationHeaderOptions.FunctionArn
+        );
 
-            clonedRequest.Content = clonedContent;
-        }
+        // Add any additional custom headers
+        foreach (var header in clientOptions.InvocationHeaderOptions.AdditionalHeaders)
+            response.Headers.Add(header.Key, header.Value);
 
-        return new InvocationCompletion { Request = clonedRequest, RequestType = requestType };
+        return response;
     }
+
+    private string GetRequestId() =>
+        Interlocked.Increment(ref _requestCounter).ToString().PadLeft(12, '0');
 }
diff --git a/src/AwsLambda.Host.Testing/ServerState.cs b/src/AwsLambda.Host.Testing/ServerState.cs
new file mode 100644
index 00000000..9e2f394e
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/ServerState.cs
@@ -0,0 +1,37 @@
+namespace AwsLambda.Host.Testing;
+
+/// <summary>
+/// Represents the lifecycle state of a LambdaTestServer.
+/// </summary>
+public enum ServerState
+{
+    /// <summary>
+    /// Server created but not started.
+    /// </summary>
+    Created,
+
+    /// <summary>
+    /// Server is starting (building host).
+    /// </summary>
+    Starting,
+
+    /// <summary>
+    /// Server is running and accepting invocations.
+    /// </summary>
+    Running,
+
+    /// <summary>
+    /// Server is stopping.
+    /// </summary>
+    Stopping,
+
+    /// <summary>
+    /// Server has stopped cleanly.
+    /// </summary>
+    Stopped,
+
+    /// <summary>
+    /// Server has been disposed.
+    /// </summary>
+    Disposed,
+}

From 8b43e3e20923144f266d8d75cf9c98a5f9d91c45 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Sun, 7 Dec 2025 17:35:16 -0500
Subject: [PATCH 05/29] fix(testing): correct exception type assertion and
 re-throw caught exception in disposal

- Updated `LambdaHostTest` to assert `AggregateException` instead of `TaskCanceledException`.
- Re-added a `throw` statement in `LambdaTestServer` to propagate caught exceptions during disposal.
---
 examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs | 2 +-
 src/AwsLambda.Host.Testing/LambdaTestServer.cs                  | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 26a01fb1..9d4dd9f0 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -111,7 +111,7 @@ await Assert.ThrowsAsync<TaskCanceledException>(() =>
 
     [Fact]
     public async Task InvokeAsync_WithZeroTimeout_CancelsInvocation() =>
-        await Assert.ThrowsAsync<TaskCanceledException>(async () =>
+        await Assert.ThrowsAsync<AggregateException>(async () =>
         {
             await using var factory = new LambdaApplicationFactory<Program>();
             await factory.Server.StartAsync(TestContext.Current.CancellationToken);
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServer.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
index 7e11ac3e..5adecd39 100644
--- a/src/AwsLambda.Host.Testing/LambdaTestServer.cs
+++ b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
@@ -227,6 +227,7 @@ public async ValueTask DisposeAsync()
                 }
                 catch
                 {
+                    throw;
                     // Best effort - continue with disposal
                 }
 

From 466da55b5629b9bf0b787b1df3cd474a115d3cfd Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Sun, 7 Dec 2025 20:19:09 -0500
Subject: [PATCH 06/29] feat(testing): enhance entry point completion handling
 and exception propagation

- Added `FailPendingInvocations` in `InvocationProcessor` to handle failed invocations gracefully.
- Updated `ConfigureHostBuilder` to accommodate `entryPointCompletion` for better fault handling.
- Enhanced `LambdaTestServer` with entry point completion monitoring and fault propagation methods.
- Introduced `_entryPointCompletionTcs` in `DeferredHostBuilder` for tracking entry point execution state.
---
 .../DeferredHostBuilder.cs                    |  8 ++
 .../Internal/InvocationProcessor.cs           | 12 +++
 .../LambdaApplicationFactory.cs               | 13 ++-
 .../LambdaTestServer.cs                       | 87 ++++++++++++++++++-
 4 files changed, 115 insertions(+), 5 deletions(-)

diff --git a/src/AwsLambda.Host.Testing/DeferredHostBuilder.cs b/src/AwsLambda.Host.Testing/DeferredHostBuilder.cs
index 7fe11895..6e3d1fe0 100644
--- a/src/AwsLambda.Host.Testing/DeferredHostBuilder.cs
+++ b/src/AwsLambda.Host.Testing/DeferredHostBuilder.cs
@@ -25,6 +25,10 @@ internal sealed class DeferredHostBuilder : IHostBuilder
         TaskCreationOptions.RunContinuationsAsynchronously
     );
 
+    private readonly TaskCompletionSource<Exception?> _entryPointCompletionTcs = new(
+        TaskCreationOptions.RunContinuationsAsynchronously
+    );
+
     private Action<IHostBuilder> _configure;
     private Func<string[], object>? _hostFactory;
 
@@ -39,6 +43,8 @@ public DeferredHostBuilder() =>
 
     public IDictionary<object, object> Properties { get; } = new Dictionary<object, object>();
 
+    public Task<Exception?> EntryPointCompletion => _entryPointCompletionTcs.Task;
+
     public IHost Build()
     {
         // Hosting configuration is being provided by args so that
@@ -116,10 +122,12 @@ public void EntryPointCompleted(Exception? exception)
         if (exception is not null)
         {
             _hostStartTcs.TrySetException(exception);
+            _entryPointCompletionTcs.TrySetResult(exception);
         }
         else
         {
             _hostStartTcs.TrySetResult();
+            _entryPointCompletionTcs.TrySetResult(null);
         }
     }
 
diff --git a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
index 3f596427..28efa4ee 100644
--- a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
+++ b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
@@ -400,4 +400,16 @@ HttpRequestMessage sourceRequest
 
         return new InvocationCompletion { Request = clonedRequest, RequestType = requestType };
     }
+
+    internal void FailPendingInvocations(Exception exception)
+    {
+        while (_queuedNextRequests.Reader.TryRead(out var queuedTransaction))
+            queuedTransaction.Fail(exception);
+
+        while (_transactionChannel.Reader.TryRead(out var transaction))
+            transaction.Fail(exception);
+
+        foreach (var pendingInvocation in _pendingInvocations.Values)
+            pendingInvocation.ResponseTcs.TrySetException(exception);
+    }
 }
diff --git a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
index 609eefe8..c44d9e59 100644
--- a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
+++ b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
@@ -197,7 +197,7 @@ private void EnsureServer()
             // to IHostBuilder.Build.
             deferredHostBuilder.SetHostFactory(factory);
 
-            ConfigureHostBuilder(deferredHostBuilder);
+            ConfigureHostBuilder(deferredHostBuilder, deferredHostBuilder.EntryPointCompletion);
             return;
         }
 
@@ -205,7 +205,10 @@ private void EnsureServer()
     }
 
     [MemberNotNull(nameof(_server))]
-    private void ConfigureHostBuilder(IHostBuilder hostBuilder)
+    private void ConfigureHostBuilder(
+        IHostBuilder hostBuilder,
+        Task<Exception?> entryPointCompletion
+    )
     {
         SetContentRoot(hostBuilder);
         _configuration(hostBuilder);
@@ -228,7 +231,11 @@ private void ConfigureHostBuilder(IHostBuilder hostBuilder)
         _host = CreateHost(hostBuilder);
 
         // Create the public server with the built (but not started) host
-        _server = new LambdaTestServer(_host, processor);
+        _server = new LambdaTestServer(
+            _host,
+            processor,
+            entryPointCompletion: entryPointCompletion
+        );
     }
 
     private void SetContentRoot(IHostBuilder builder)
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServer.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
index 5adecd39..db84bcc0 100644
--- a/src/AwsLambda.Host.Testing/LambdaTestServer.cs
+++ b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
@@ -16,20 +16,24 @@ public class LambdaTestServer : IAsyncDisposable
     private readonly InvocationProcessor _processor;
     private readonly LambdaClientOptions _defaultClientOptions;
     private readonly JsonSerializerOptions _jsonSerializerOptions;
+    private readonly Task<Exception?> _entryPointCompletion;
     private int _requestCounter;
+    private Task? _entryPointMonitorTask;
     private ServerState _state;
 
     internal LambdaTestServer(
         IHost host,
         InvocationProcessor processor,
         LambdaClientOptions? defaultClientOptions = null,
-        JsonSerializerOptions? jsonSerializerOptions = null
+        JsonSerializerOptions? jsonSerializerOptions = null,
+        Task<Exception?>? entryPointCompletion = null
     )
     {
         _host = host ?? throw new ArgumentNullException(nameof(host));
         _processor = processor ?? throw new ArgumentNullException(nameof(processor));
         _defaultClientOptions = defaultClientOptions ?? new LambdaClientOptions();
         _jsonSerializerOptions = jsonSerializerOptions ?? new JsonSerializerOptions();
+        _entryPointCompletion = entryPointCompletion ?? Task.FromResult<Exception?>(null);
         _state = ServerState.Created;
     }
 
@@ -72,6 +76,8 @@ public async Task StartAsync(CancellationToken cancellationToken = default)
         {
             _state = ServerState.Starting;
 
+            EnsureEntryPointDidNotFault();
+
             // Start the host
             await _host.StartAsync(cancellationToken);
 
@@ -79,6 +85,7 @@ public async Task StartAsync(CancellationToken cancellationToken = default)
             _processor.StartProcessing();
 
             _state = ServerState.Running;
+            BeginEntryPointMonitoring();
         }
         catch
         {
@@ -227,10 +234,19 @@ public async ValueTask DisposeAsync()
                 }
                 catch
                 {
-                    throw;
                     // Best effort - continue with disposal
                 }
 
+            if (_entryPointMonitorTask != null)
+                try
+                {
+                    await _entryPointMonitorTask.ConfigureAwait(false);
+                }
+                catch
+                {
+                    // Swallow monitoring errors during disposal
+                }
+
             // Dispose processor
             await _processor.DisposeAsync();
 
@@ -292,4 +308,71 @@ LambdaClientOptions clientOptions
 
     private string GetRequestId() =>
         Interlocked.Increment(ref _requestCounter).ToString().PadLeft(12, '0');
+
+    private void EnsureEntryPointDidNotFault()
+    {
+        if (!_entryPointCompletion.IsCompleted)
+            return;
+
+        PropagateEntryPointCompletion(_entryPointCompletion);
+    }
+
+    private void BeginEntryPointMonitoring()
+    {
+        if (_entryPointCompletion.IsCompleted)
+        {
+            PropagateEntryPointCompletion(_entryPointCompletion);
+            return;
+        }
+
+        _entryPointMonitorTask = MonitorEntryPointCompletionAsync();
+    }
+
+    private async Task MonitorEntryPointCompletionAsync()
+    {
+        try
+        {
+            var exception = await _entryPointCompletion.ConfigureAwait(false);
+            if (exception != null)
+                await HandleEntryPointFailureAsync(exception).ConfigureAwait(false);
+        }
+        catch (Exception ex)
+        {
+            await HandleEntryPointFailureAsync(ex).ConfigureAwait(false);
+        }
+    }
+
+    private void PropagateEntryPointCompletion(Task<Exception?> completionTask)
+    {
+        if (completionTask.IsFaulted)
+            throw completionTask.Exception!.GetBaseException();
+
+        if (completionTask.Status == TaskStatus.Canceled)
+            throw new TaskCanceledException("Entry point execution was canceled.");
+
+        var exception =
+            completionTask.Status == TaskStatus.RanToCompletion ? completionTask.Result : null;
+
+        if (exception != null)
+            throw exception;
+    }
+
+    private async Task HandleEntryPointFailureAsync(Exception exception)
+    {
+        if (_state == ServerState.Disposed)
+            return;
+
+        _processor.FailPendingInvocations(exception);
+
+        try
+        {
+            await _host.StopAsync();
+        }
+        catch
+        {
+            // Best effort stop
+        }
+
+        _state = ServerState.Stopped;
+    }
 }

From 0734253e5581e98729877299134d6c502b87fdce Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Mon, 8 Dec 2025 18:12:14 -0500
Subject: [PATCH 07/29] feat(testing): enhance lifecycle handling and
 initialization in InvocationProcessor and LambdaTestServer

- Introduced `ProcessorState` enum to track `InvocationProcessor` lifecycle states.
- Enhanced `InvocationProcessor` with initialization error handling and server state transitions.
- Added `HandlePostInitErrorAsync` in `InvocationProcessor` to process initialization errors.
- Updated `LambdaTestServer.StartAsync` to wait for or fail on host initialization.
- Created `InitResponse` model to encapsulate initialization result in testing framework.
- Updated `runtime-api.yaml` with new POST `/runtime/init/error` endpoint specification.
- Extended `RequestType` and `LambdaRuntimeRouteManager` to support `PostInitError` requests.
---
 .../Lambda/Program.cs                         |   5 +
 .../Tests/LambdaHostTest.cs                   |  25 +-
 .../Internal/InvocationProcessor.cs           | 192 ++++++++++----
 .../Internal/ProcessorState.cs                |  35 +++
 .../LambdaRuntimeRouteManager.cs              |   9 +
 .../LambdaTestServer.cs                       | 126 +++------
 .../Models/ErrorResponse.cs                   |   3 +
 .../Models/InitResponse.cs                    |   7 +
 .../Models/RequestType.cs                     |   1 +
 src/AwsLambda.Host.Testing/README.md          |  31 +++
 src/AwsLambda.Host.Testing/runtime-api.yaml   | 250 ++++++++++++++++++
 11 files changed, 531 insertions(+), 153 deletions(-)
 create mode 100644 src/AwsLambda.Host.Testing/Internal/ProcessorState.cs
 create mode 100644 src/AwsLambda.Host.Testing/Models/InitResponse.cs
 create mode 100644 src/AwsLambda.Host.Testing/runtime-api.yaml

diff --git a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
index fc3329b0..b5fa7ee5 100644
--- a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
@@ -13,6 +13,11 @@
 // Build the Lambda application
 var lambda = builder.Build();
 
+lambda.OnInit(() =>
+{
+    return false;
+});
+
 // Map your handler - the event is automatically injected
 lambda.MapHandler(([Event] string name) => $"Hello {name}!");
 
diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 9d4dd9f0..7b400362 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -13,7 +13,13 @@ public class LambdaHostTest
     public async Task LambdaHost_CanStartWithoutError()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        if (!setup.InitSuccess)
+            throw new Exception("Failed to initialize Lambda host");
+
+        var x = 0;
 
         var response = await factory.Server.InvokeAsync<string, string>(
             "Jonas",
@@ -111,7 +117,7 @@ await Assert.ThrowsAsync<TaskCanceledException>(() =>
 
     [Fact]
     public async Task InvokeAsync_WithZeroTimeout_CancelsInvocation() =>
-        await Assert.ThrowsAsync<AggregateException>(async () =>
+        await Assert.ThrowsAsync<TaskCanceledException>(async () =>
         {
             await using var factory = new LambdaApplicationFactory<Program>();
             await factory.Server.StartAsync(TestContext.Current.CancellationToken);
@@ -125,4 +131,19 @@ await factory.Server.InvokeAsync<string, string>(
                 TestContext.Current.CancellationToken
             );
         });
+
+    [Fact]
+    public async Task StartAsync_WithFailingInit_ReturnsInitError()
+    {
+        // This test verifies that when OnInit returns false (as configured in Program.cs),
+        // the runtime posts to /runtime/init/error and StartAsync returns InitResponse with error
+        await using var factory = new LambdaApplicationFactory<Program>();
+
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        var initResponse = await factory.Server.StartAsync(cts.Token);
+
+        Assert.False(initResponse.InitSuccess);
+        Assert.NotNull(initResponse.Error);
+        Assert.Equal(ServerState.Stopped, factory.Server.State);
+    }
 }
diff --git a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
index 28efa4ee..7e4f503d 100644
--- a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
+++ b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
@@ -1,5 +1,6 @@
 using System.Collections.Concurrent;
 using System.Net;
+using System.Net.Http.Json;
 using System.Text;
 using System.Text.Json;
 using System.Threading.Channels;
@@ -14,19 +15,16 @@ namespace AwsLambda.Host.Testing;
 internal class InvocationProcessor : IAsyncDisposable
 {
     private const int TransactionChannelCapacity = 1024;
-    private const int NextRequestChannelCapacity = 1024;
-
-    private static readonly OperationCanceledException DisposedException = new(
-        "InvocationProcessor disposed"
-    );
 
     private readonly JsonSerializerOptions _jsonSerializerOptions;
     private readonly ConcurrentQueue<string> _pendingInvocationIds;
     private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
-    private readonly Channel<LambdaHttpTransaction> _queuedNextRequests;
     private readonly ILambdaRuntimeRouteManager _routeManager;
     private readonly CancellationTokenSource _shutdownCts;
     private readonly Channel<LambdaHttpTransaction> _transactionChannel;
+    private readonly TaskCompletionSource<InitResponse> _initCompletionTcs;
+    private readonly SemaphoreSlim _invocationAddedSignal;
+    private ProcessorState _state;
     private Task? _processingTask;
 
     internal InvocationProcessor(
@@ -44,33 +42,34 @@ internal InvocationProcessor(
         );
         _pendingInvocationIds = new ConcurrentQueue<string>();
         _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
-        _queuedNextRequests = Channel.CreateBounded<LambdaHttpTransaction>(
-            new BoundedChannelOptions(NextRequestChannelCapacity)
-            {
-                SingleReader = false,
-                SingleWriter = false,
-                FullMode = BoundedChannelFullMode.Wait,
-            }
-        );
         _routeManager = routeManager ?? new LambdaRuntimeRouteManager();
         _jsonSerializerOptions = jsonSerializerOptions ?? new JsonSerializerOptions();
         _shutdownCts = new CancellationTokenSource();
+        _initCompletionTcs = new TaskCompletionSource<InitResponse>(
+            TaskCreationOptions.RunContinuationsAsynchronously
+        );
+        _invocationAddedSignal = new SemaphoreSlim(0);
+        _state = ProcessorState.Created;
     }
 
     public async ValueTask DisposeAsync()
     {
+        // Transition to stopping state
+        if (_state != ProcessorState.Stopped)
+            _state = ProcessorState.Stopping;
+
         await _shutdownCts.CancelAsync();
 
-        _transactionChannel.Writer.TryComplete();
-        _queuedNextRequests.Writer.TryComplete();
+        // Complete init if still pending (server shutting down before init)
+        if (_state == ProcessorState.Initializing || _state == ProcessorState.Stopping)
+            _initCompletionTcs.TrySetCanceled(_shutdownCts.Token);
 
-        // Cancel any transactions waiting for work
-        while (_queuedNextRequests.Reader.TryRead(out var queuedTransaction))
-            queuedTransaction.Fail(DisposedException);
+        _transactionChannel.Writer.TryComplete();
 
         // Fail any in-flight transactions that haven't been processed yet
+        var disposedException = new OperationCanceledException("InvocationProcessor disposed");
         while (_transactionChannel.Reader.TryRead(out var transaction))
-            transaction.Fail(DisposedException);
+            transaction.Fail(disposedException);
 
         // Cancel any pending invocations waiting for bootstrap responses
         foreach (var pendingInvocation in _pendingInvocations.Values)
@@ -86,9 +85,18 @@ public async ValueTask DisposeAsync()
                 // Expected when task is canceled
             }
 
+        _state = ProcessorState.Stopped;
         _shutdownCts.Dispose();
+        _invocationAddedSignal.Dispose();
     }
 
+    /// <summary>
+    /// Gets a task that completes when initialization finishes (either success or failure).
+    /// Used by LambdaTestServer.StartAsync to wait for initialization to complete.
+    /// </summary>
+    /// <returns>Task that completes with InitResponse indicating success or failure.</returns>
+    internal Task<InitResponse> GetInitCompletionAsync() => _initCompletionTcs.Task;
+
     /// <summary>
     /// Creates the HTTP handler for Lambda Bootstrap to use.
     /// </summary>
@@ -104,6 +112,7 @@ internal void StartProcessing()
         if (_processingTask != null)
             throw new InvalidOperationException("Processor already started");
 
+        _state = ProcessorState.Initializing;
         _processingTask = Task.Run(ProcessTransactionsAsync);
     }
 
@@ -118,6 +127,11 @@ internal async Task<InvocationCompletion> QueueInvocationAsync(
         CancellationToken cancellationToken
     )
     {
+        if (_state != ProcessorState.Running)
+            throw new InvalidOperationException(
+                $"Cannot queue invocation when processor is in {_state} state. Processor must be in Running state."
+            );
+
         var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
 
         if (!_pendingInvocations.TryAdd(requestId, pending))
@@ -125,9 +139,8 @@ CancellationToken cancellationToken
 
         _pendingInvocationIds.Enqueue(requestId);
 
-        // If there's a queued /next request, serve it immediately
-        if (_queuedNextRequests.Reader.TryRead(out var nextTransaction))
-            RespondToNextRequest(nextTransaction);
+        // Signal any waiting /next request that work is available
+        _invocationAddedSignal.Release();
 
         using var cancellationRegistration = cancellationToken.Register(() =>
             CancelPendingInvocation(requestId, cancellationToken)
@@ -180,6 +193,10 @@ private async Task ProcessTransactionAsync(LambdaHttpTransaction transaction)
                 await HandlePostErrorAsync(transaction, routeValues!);
                 break;
 
+            case RequestType.PostInitError:
+                await HandlePostInitErrorAsync(transaction);
+                break;
+
             default:
                 throw new InvalidOperationException(
                     $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
@@ -192,23 +209,49 @@ private async Task ProcessTransactionAsync(LambdaHttpTransaction transaction)
     /// </summary>
     private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)
     {
-        // Try to dequeue next pending invocation (FIFO)
-        if (!TryDequeuePendingInvocation(out var pending))
+        // First successful /next call indicates initialization succeeded
+        if (_state == ProcessorState.Initializing)
+        {
+            _state = ProcessorState.Running;
+            _initCompletionTcs.TrySetResult(new InitResponse { InitSuccess = true });
+        }
+
+        // Loop until we find work or shutdown
+        while (true)
         {
-            // No work available - queue this /next request
+            // Check if processor is still in a valid state to serve /next requests
+            if (_state != ProcessorState.Running && _state != ProcessorState.Initializing)
+            {
+                transaction.Fail(
+                    new InvalidOperationException(
+                        $"Processor is in {_state} state, cannot serve /next request"
+                    )
+                );
+                return;
+            }
+
+            // Try to dequeue next pending invocation (FIFO)
+            if (TryDequeuePendingInvocation(out var pending))
+            {
+                RespondToNextRequest(transaction, pending);
+                return;
+            }
+
+            // No work available - wait for new invocation or shutdown
             try
             {
-                await _queuedNextRequests.Writer.WriteAsync(transaction, _shutdownCts.Token);
+                await _invocationAddedSignal.WaitAsync(_shutdownCts.Token);
+                // Loop back to check state and try dequeue again
             }
             catch (OperationCanceledException)
             {
-                transaction.Fail(DisposedException);
+                // Server shutting down
+                transaction.Fail(
+                    new OperationCanceledException("Server stopped while waiting for work")
+                );
+                return;
             }
-
-            return;
         }
-
-        RespondToNextRequest(transaction, pending);
     }
 
     /// <summary>
@@ -224,27 +267,6 @@ private void RespondToNextRequest(LambdaHttpTransaction transaction, PendingInvo
         _pendingInvocationIds.Enqueue(pending.RequestId);
     }
 
-    /// <summary>
-    /// Responds to a /next request with a pending invocation by looking up the next one.
-    /// </summary>
-    private void RespondToNextRequest(LambdaHttpTransaction transaction)
-    {
-        // Try to dequeue next pending invocation (FIFO)
-        if (!TryDequeuePendingInvocation(out var pending))
-        {
-            // This shouldn't happen, but if it does, respond with error
-            transaction.Respond(
-                new HttpResponseMessage(HttpStatusCode.InternalServerError)
-                {
-                    Content = new StringContent("No pending invocations available"),
-                }
-            );
-            return;
-        }
-
-        RespondToNextRequest(transaction, pending);
-    }
-
     /// <summary>
     /// Handles POST /invocation/{requestId}/response - successful function execution.
     /// </summary>
@@ -342,6 +364,67 @@ await CreateCompletionAsync(RequestType.PostError, transaction.Request)
         );
     }
 
+    /// <summary>
+    /// Handles POST /runtime/init/error - Lambda initialization failed.
+    /// Captures the error details and marks initialization as failed.
+    /// The server will be automatically stopped after this is reported.
+    /// </summary>
+    private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction)
+    {
+        // Parse error details from request body
+        ErrorResponse? errorResponse = null;
+        if (transaction.Request.Content != null)
+            try
+            {
+                errorResponse = await transaction.Request.Content.ReadFromJsonAsync<ErrorResponse>(
+                    _jsonSerializerOptions
+                );
+            }
+            catch
+            {
+                // Fallback: create basic error from body text
+                var body = await transaction.Request.Content.ReadAsStringAsync();
+                errorResponse = new ErrorResponse
+                {
+                    ErrorType = transaction.Request.Headers.TryGetValues(
+                        "Lambda-Runtime-Function-Error-Type",
+                        out var values
+                    )
+                        ? values.FirstOrDefault() ?? "Unknown"
+                        : "Unknown",
+                    ErrorMessage = body,
+                };
+            }
+
+        // Mark initialization as failed (only once)
+        if (_state == ProcessorState.Initializing)
+        {
+            _state = ProcessorState.Stopped;
+            _initCompletionTcs.TrySetResult(
+                new InitResponse { InitSuccess = false, Error = errorResponse }
+            );
+            // Wake up any waiting /next requests by releasing the semaphore
+            // They will check state and fail appropriately
+            _invocationAddedSignal.Release();
+        }
+
+        // Acknowledge to Bootstrap with 202 Accepted
+        transaction.Respond(
+            new HttpResponseMessage(HttpStatusCode.Accepted)
+            {
+                Content = new StringContent(
+                    JsonSerializer.Serialize(
+                        new Dictionary<string, string> { ["status"] = "OK" },
+                        _jsonSerializerOptions
+                    ),
+                    Encoding.UTF8,
+                    "application/json"
+                ),
+                Version = Version.Parse("1.1"),
+            }
+        );
+    }
+
     private bool TryDequeuePendingInvocation(out PendingInvocation pendingInvocation)
     {
         var now = DateTimeOffset.UtcNow;
@@ -403,9 +486,6 @@ HttpRequestMessage sourceRequest
 
     internal void FailPendingInvocations(Exception exception)
     {
-        while (_queuedNextRequests.Reader.TryRead(out var queuedTransaction))
-            queuedTransaction.Fail(exception);
-
         while (_transactionChannel.Reader.TryRead(out var transaction))
             transaction.Fail(exception);
 
diff --git a/src/AwsLambda.Host.Testing/Internal/ProcessorState.cs b/src/AwsLambda.Host.Testing/Internal/ProcessorState.cs
new file mode 100644
index 00000000..35ad2269
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/Internal/ProcessorState.cs
@@ -0,0 +1,35 @@
+namespace AwsLambda.Host.Testing;
+
+/// <summary>
+/// Represents the lifecycle state of the InvocationProcessor.
+/// </summary>
+internal enum ProcessorState
+{
+    /// <summary>
+    /// Processor has been created but not yet started.
+    /// </summary>
+    Created,
+
+    /// <summary>
+    /// Processor is running and waiting for initialization to complete.
+    /// Transitions to Running on first successful /next request or to Stopped on init error.
+    /// </summary>
+    Initializing,
+
+    /// <summary>
+    /// Processor is running and processing invocations.
+    /// Initialization completed successfully.
+    /// </summary>
+    Running,
+
+    /// <summary>
+    /// Processor is shutting down.
+    /// Transitions to Stopped when shutdown completes.
+    /// </summary>
+    Stopping,
+
+    /// <summary>
+    /// Processor has stopped and is no longer processing requests.
+    /// </summary>
+    Stopped,
+}
diff --git a/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs b/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs
index e60d02c3..4ce26746 100644
--- a/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs
+++ b/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs
@@ -18,6 +18,15 @@ internal class LambdaRuntimeRouteManager : ILambdaRuntimeRouteManager
             ),
         },
         new()
+        {
+            Type = RequestType.PostInitError,
+            Method = HttpMethod.Post.Method,
+            Matcher = new TemplateMatcher(
+                TemplateParser.Parse("{version}/runtime/init/error"),
+                new RouteValueDictionary()
+            ),
+        },
+        new()
         {
             Type = RequestType.PostResponse,
             Method = HttpMethod.Post.Method,
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServer.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
index db84bcc0..2d8048ad 100644
--- a/src/AwsLambda.Host.Testing/LambdaTestServer.cs
+++ b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
@@ -18,7 +18,6 @@ public class LambdaTestServer : IAsyncDisposable
     private readonly JsonSerializerOptions _jsonSerializerOptions;
     private readonly Task<Exception?> _entryPointCompletion;
     private int _requestCounter;
-    private Task? _entryPointMonitorTask;
     private ServerState _state;
 
     internal LambdaTestServer(
@@ -47,25 +46,17 @@ internal LambdaTestServer(
     /// Only available after the server has been started.
     /// </summary>
     /// <exception cref="InvalidOperationException">Thrown if server is not in Running state.</exception>
-    public IServiceProvider Services
-    {
-        get
-        {
-            if (_state != ServerState.Running)
-                throw new InvalidOperationException(
-                    $"Services are only available when server is Running. Current state: {_state}. Call StartAsync() first."
-                );
-
-            return _host.Services;
-        }
-    }
+    public IServiceProvider Services => _host.Services;
 
     /// <summary>
-    /// Starts the Lambda host and begins processing invocations.
+    /// Starts the Lambda host and waits for initialization to complete.
+    /// Returns initialization result indicating success or failure.
+    /// If initialization fails, the server is automatically stopped.
     /// </summary>
     /// <param name="cancellationToken">Cancellation token.</param>
+    /// <returns>InitResponse indicating whether initialization succeeded or failed.</returns>
     /// <exception cref="InvalidOperationException">Thrown if server is not in Created state.</exception>
-    public async Task StartAsync(CancellationToken cancellationToken = default)
+    public async Task<InitResponse> StartAsync(CancellationToken cancellationToken = default)
     {
         if (_state != ServerState.Created)
             throw new InvalidOperationException(
@@ -76,8 +67,6 @@ public async Task StartAsync(CancellationToken cancellationToken = default)
         {
             _state = ServerState.Starting;
 
-            EnsureEntryPointDidNotFault();
-
             // Start the host
             await _host.StartAsync(cancellationToken);
 
@@ -85,7 +74,31 @@ public async Task StartAsync(CancellationToken cancellationToken = default)
             _processor.StartProcessing();
 
             _state = ServerState.Running;
-            BeginEntryPointMonitoring();
+
+            // Race initialization completion against entry point failure
+            var initTask = _processor.GetInitCompletionAsync();
+            var completed = await Task.WhenAny(initTask, _entryPointCompletion)
+                .WaitAsync(cancellationToken);
+
+            // Check if entry point failed before init completed
+            if (completed == _entryPointCompletion)
+            {
+                var exception = await _entryPointCompletion;
+                if (exception != null)
+                    throw new InvalidOperationException(
+                        "Entry point failed during initialization",
+                        exception
+                    );
+            }
+
+            // Wait for init to complete
+            var initResponse = await initTask;
+
+            // If init failed, auto-stop the server
+            if (!initResponse.InitSuccess)
+                await StopAsync(CancellationToken.None);
+
+            return initResponse;
         }
         catch
         {
@@ -237,16 +250,6 @@ public async ValueTask DisposeAsync()
                     // Best effort - continue with disposal
                 }
 
-            if (_entryPointMonitorTask != null)
-                try
-                {
-                    await _entryPointMonitorTask.ConfigureAwait(false);
-                }
-                catch
-                {
-                    // Swallow monitoring errors during disposal
-                }
-
             // Dispose processor
             await _processor.DisposeAsync();
 
@@ -308,71 +311,4 @@ LambdaClientOptions clientOptions
 
     private string GetRequestId() =>
         Interlocked.Increment(ref _requestCounter).ToString().PadLeft(12, '0');
-
-    private void EnsureEntryPointDidNotFault()
-    {
-        if (!_entryPointCompletion.IsCompleted)
-            return;
-
-        PropagateEntryPointCompletion(_entryPointCompletion);
-    }
-
-    private void BeginEntryPointMonitoring()
-    {
-        if (_entryPointCompletion.IsCompleted)
-        {
-            PropagateEntryPointCompletion(_entryPointCompletion);
-            return;
-        }
-
-        _entryPointMonitorTask = MonitorEntryPointCompletionAsync();
-    }
-
-    private async Task MonitorEntryPointCompletionAsync()
-    {
-        try
-        {
-            var exception = await _entryPointCompletion.ConfigureAwait(false);
-            if (exception != null)
-                await HandleEntryPointFailureAsync(exception).ConfigureAwait(false);
-        }
-        catch (Exception ex)
-        {
-            await HandleEntryPointFailureAsync(ex).ConfigureAwait(false);
-        }
-    }
-
-    private void PropagateEntryPointCompletion(Task<Exception?> completionTask)
-    {
-        if (completionTask.IsFaulted)
-            throw completionTask.Exception!.GetBaseException();
-
-        if (completionTask.Status == TaskStatus.Canceled)
-            throw new TaskCanceledException("Entry point execution was canceled.");
-
-        var exception =
-            completionTask.Status == TaskStatus.RanToCompletion ? completionTask.Result : null;
-
-        if (exception != null)
-            throw exception;
-    }
-
-    private async Task HandleEntryPointFailureAsync(Exception exception)
-    {
-        if (_state == ServerState.Disposed)
-            return;
-
-        _processor.FailPendingInvocations(exception);
-
-        try
-        {
-            await _host.StopAsync();
-        }
-        catch
-        {
-            // Best effort stop
-        }
-
-        _state = ServerState.Stopped;
-    }
 }
diff --git a/src/AwsLambda.Host.Testing/Models/ErrorResponse.cs b/src/AwsLambda.Host.Testing/Models/ErrorResponse.cs
index 540fd345..02fb088c 100644
--- a/src/AwsLambda.Host.Testing/Models/ErrorResponse.cs
+++ b/src/AwsLambda.Host.Testing/Models/ErrorResponse.cs
@@ -13,6 +13,9 @@ public class ErrorResponse
     [JsonPropertyName("cause")]
     public ErrorCause? Cause { get; set; }
 
+    [JsonPropertyName("causes")]
+    public List<ErrorCause>? Causes { get; set; }
+
     /// <summary>
     /// The error message describing what went wrong.
     /// </summary>
diff --git a/src/AwsLambda.Host.Testing/Models/InitResponse.cs b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
new file mode 100644
index 00000000..d861159a
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
@@ -0,0 +1,7 @@
+namespace AwsLambda.Host.Testing;
+
+public class InitResponse
+{
+    public ErrorResponse? Error { get; internal set; }
+    public bool InitSuccess { get; internal set; }
+}
diff --git a/src/AwsLambda.Host.Testing/Models/RequestType.cs b/src/AwsLambda.Host.Testing/Models/RequestType.cs
index 935bd790..49108da9 100644
--- a/src/AwsLambda.Host.Testing/Models/RequestType.cs
+++ b/src/AwsLambda.Host.Testing/Models/RequestType.cs
@@ -5,4 +5,5 @@ internal enum RequestType
     GetNextInvocation,
     PostResponse,
     PostError,
+    PostInitError,
 }
diff --git a/src/AwsLambda.Host.Testing/README.md b/src/AwsLambda.Host.Testing/README.md
index afb866ba..547ccc1c 100644
--- a/src/AwsLambda.Host.Testing/README.md
+++ b/src/AwsLambda.Host.Testing/README.md
@@ -178,6 +178,37 @@ Headers:
 
 ```
 
+```json
+{
+  "errorType": "AggregateException",
+  "errorMessage": "Encountered errors while running OnInit handlers: (Init failed)",
+  "stackTrace": [
+    "at AwsLambda.Host.Builder.LambdaOnInitBuilder.<Build>b__10_1(CancellationToken stoppingToken) in /Users/jonasha/Repos/CSharp/dotnet-lambda-host/src/AwsLambda.Host/Builder/LambdaOnInitBuilder.cs:line 65",
+    "at Amazon.Lambda.RuntimeSupport.LambdaBootstrap.InitializeAsync()"
+  ],
+  "cause": {
+    "errorType": "Exception",
+    "errorMessage": "Init failed",
+    "stackTrace": [
+      "at Program.<>c.<<Main>$>b__0_1() in /Users/jonasha/Repos/CSharp/dotnet-lambda-host/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs:line 18",
+      "at AwsLambda.Host.Core.Generated.<LambdaHandler_g>F0316F908774B11B609C941996D283BB313AD08F671E39F0BE1901FEE2BBA6D4C__GeneratedLambdaOnInitBuilderExtensions.<>c__DisplayClass0_0.<OnInitInterceptor0>g__OnInit|0(IServiceProvider serviceProvider, CancellationToken cancellationToken) in /Users/jonasha/Repos/CSharp/dotnet-lambda-host/examples/AwsLambda.Host.Example.Testing/Lambda/obj/Debug/net10.0/AwsLambda.Host.SourceGenerators/AwsLambda.Host.SourceGenerators.MapHandlerIncrementalGenerator/LambdaHandler.g.cs:line 109",
+      "at AwsLambda.Host.Builder.LambdaOnInitBuilder.RunInitHandler(LambdaInitDelegate handler, CancellationToken cancellationToken) in /Users/jonasha/Repos/CSharp/dotnet-lambda-host/src/AwsLambda.Host/Builder/LambdaOnInitBuilder.cs:line 81"
+    ]
+  },
+  "causes": [
+    {
+      "errorType": "Exception",
+      "errorMessage": "Init failed",
+      "stackTrace": [
+        "at Program.<>c.<<Main>$>b__0_1() in /Users/jonasha/Repos/CSharp/dotnet-lambda-host/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs:line 18",
+        "at AwsLambda.Host.Core.Generated.<LambdaHandler_g>F0316F908774B11B609C941996D283BB313AD08F671E39F0BE1901FEE2BBA6D4C__GeneratedLambdaOnInitBuilderExtensions.<>c__DisplayClass0_0.<OnInitInterceptor0>g__OnInit|0(IServiceProvider serviceProvider, CancellationToken cancellationToken) in /Users/jonasha/Repos/CSharp/dotnet-lambda-host/examples/AwsLambda.Host.Example.Testing/Lambda/obj/Debug/net10.0/AwsLambda.Host.SourceGenerators/AwsLambda.Host.SourceGenerators.MapHandlerIncrementalGenerator/LambdaHandler.g.cs:line 109",
+        "at AwsLambda.Host.Builder.LambdaOnInitBuilder.RunInitHandler(LambdaInitDelegate handler, CancellationToken cancellationToken) in /Users/jonasha/Repos/CSharp/dotnet-lambda-host/src/AwsLambda.Host/Builder/LambdaOnInitBuilder.cs:line 81"
+      ]
+    }
+  ]
+}
+```
+
 # In-Memory Lambda Testing Client Implementation Summary
 
 ## Overview
diff --git a/src/AwsLambda.Host.Testing/runtime-api.yaml b/src/AwsLambda.Host.Testing/runtime-api.yaml
new file mode 100644
index 00000000..9acfba2d
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/runtime-api.yaml
@@ -0,0 +1,250 @@
+#
+# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#  http://aws.amazon.com/apache2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+#
+# This file was copied from https://docs.aws.amazon.com/lambda/latest/dg/samples/runtime-api.zip on 12/1/2018
+#
+# This file documents BYOL (RAPID) API using Open API 3 specification format.
+#
+#
+# A note on error reporting in BYOL.
+#
+# Runtimes should report all errors using Lambda standard error format, in order to integrate with other AWS services:
+#
+# Content-Type: application/vnd.aws.lambda.error+json:
+# {
+#     "errorMessage": "...",
+#     "errorType": "...",
+#     "stackTrace": [],
+# }
+#
+# Corresponding Open API 3 schema:
+#
+# ErrorRequest:
+#   type: object
+#   properties:
+#     errorMessage:
+#       type: string
+#     errorType:
+#       type: string
+#     stackTrace:
+#       type: array
+#       items:
+#         type: string
+#
+# Lambda's default behavior is to use Lambda-Runtime-Function-Error-Type header value to construct an error response
+# when error payload is not provided or can not be read.
+
+
+openapi: 3.0.0
+info:
+  title: BYOL API
+  description: Native Runtime API.
+  version: 1.0.2
+
+servers:
+  - url: /2018-06-01
+
+paths:
+
+  /runtime/init/error:
+    post:
+      summary: >
+        Non-recoverable initialization error. Runtime should exit after reporting
+        the error. Error will be served in response to the first invoke.
+      parameters:
+        - in: header
+          name: Lambda-Runtime-Function-Error-Type
+          schema:
+            type: string
+      requestBody:
+        content:
+          '*/*':
+            schema: { }
+      responses:
+        '202':
+          description: Accepted
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/StatusResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Runtime should exit promptly.
+
+  /runtime/invocation/next:
+    get:
+      summary: >
+        Runtime makes this HTTP request when it is ready to receive and process a
+        new invoke.
+      responses:
+        '200':
+          description: >
+            This is an iterator-style blocking API call. Response contains
+            event JSON document, specific to the invoking service.
+          headers:
+            Lambda-Runtime-Aws-Request-Id:
+              description: AWS request ID associated with the request.
+              schema:
+                type: string
+            Lambda-Runtime-Trace-Id:
+              description: X-Ray tracing header.
+              schema:
+                type: string
+            Lambda-Runtime-Client-Context:
+              description: >
+                Information about the client application and device when invoked
+                through the AWS Mobile SDK.
+              schema:
+                type: string
+            Lambda-Runtime-Cognito-Identity:
+              description: >
+                Information about the Amazon Cognito identity provider when invoked
+                through the AWS Mobile SDK.
+              schema:
+                type: string
+            Lambda-Runtime-Deadline-Ms:
+              description: >
+                Function execution deadline counted in milliseconds since the Unix epoch.
+              schema:
+                type: string
+            Lambda-Runtime-Invoked-Function-Arn:
+              description: >
+                The ARN requested. This can be different in each invoke that
+                executes the same version.
+              schema:
+                type: string
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/EventResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Runtime should exit promptly.
+
+  /runtime/invocation/{AwsRequestId}/response:
+    post:
+      summary: Runtime makes this request in order to submit a response.
+      parameters:
+        - in: path
+          name: AwsRequestId
+          schema:
+            type: string
+          required: true
+      requestBody:
+        content:
+          '*/*':
+            schema: { }
+      responses:
+        '202':
+          description: Accepted
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/StatusResponse'
+        '400':
+          description: Bad Request
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '413':
+          description: Payload Too Large
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Runtime should exit promptly.
+
+  /runtime/invocation/{AwsRequestId}/error:
+    post:
+      summary: >
+        Runtime makes this request in order to submit an error response. It can
+        be either a function error, or a runtime error. Error will be served in
+        response to the invoke.
+      parameters:
+        - in: path
+          name: AwsRequestId
+          schema:
+            type: string
+          required: true
+        - in: header
+          name: Lambda-Runtime-Function-Error-Type
+          schema:
+            type: string
+      requestBody:
+        content:
+          '*/*':
+            schema: { }
+      responses:
+        '202':
+          description: Accepted
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/StatusResponse'
+        '400':
+          description: Bad Request
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Runtime should exit promptly.
+
+components:
+  schemas:
+
+    StatusResponse:
+      type: object
+      properties:
+        status:
+          type: string
+
+    ErrorResponse:
+      type: object
+      properties:
+        errorMessage:
+          type: string
+        errorType:
+          type: string
+
+    EventResponse:
+      type: object
\ No newline at end of file

From df7c446367cc010cbe0c3bd46a59efa750b51fbe Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Mon, 8 Dec 2025 18:32:22 -0500
Subject: [PATCH 08/29] feat(testing): add LambdaServerV2 for async server
 lifecycle management

- Implemented `LambdaServerV2` class with async start, invoke, and stop methods.
- Added `DisposeAsync` method for proper resource cleanup.
- Introduced `InvokeAsync` operation with generic request/response handling.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 57 ++++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 src/AwsLambda.Host.Testing/LambdaServerV2.cs

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
new file mode 100644
index 00000000..651c953b
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -0,0 +1,57 @@
+namespace AwsLambda.Host.Testing;
+
+public class LambdaServerV2 : IAsyncDisposable
+{
+    private bool _started;
+    private bool _stopped;
+
+    public async ValueTask DisposeAsync()
+    {
+        if (!_stopped)
+            await StopAsync();
+
+        // TODO release managed resources here
+    }
+
+    public static LambdaServerV2 Create() => new();
+
+    //      ┌──────────────────────────────────────────────────────────┐
+    //      │                        Public API                        │
+    //      └──────────────────────────────────────────────────────────┘
+
+    public async Task StartAsync(CancellationToken cancellationToken = default) => _started = true;
+
+    public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
+        TEvent invokeEvent,
+        CancellationToken cancellationToken = default
+    )
+    {
+        if (!_started)
+            throw new InvalidOperationException("Server is not started.");
+
+        return default;
+    }
+
+    public async Task StopAsync(CancellationToken cancellationToken = default)
+    {
+        if (!_started)
+            return;
+
+        _stopped = true;
+    }
+
+    //      ┌──────────────────────────────────────────────────────────┐
+    //      │                  Internal Server Logic                   │
+    //      └──────────────────────────────────────────────────────────┘
+}
+
+public static class Temp
+{
+    public static async Task Run()
+    {
+        await using var server = LambdaServerV2.Create();
+        await server.StartAsync();
+        var result = await server.InvokeAsync<string, string>("Jonas", CancellationToken.None);
+        await server.StopAsync();
+    }
+}

From dca6cbb50a7972adab77c0af9820857420febc78 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Mon, 8 Dec 2025 19:06:06 -0500
Subject: [PATCH 09/29] feat(abstractions): introduce
 DefaultLambdaJsonSerializerOptions for standardized JSON setup

- Added `DefaultLambdaJsonSerializerOptions` to centralize JSON serializer configuration for AWS Lambda.
- Updated `EnvelopeOptions` to use `DefaultLambdaJsonSerializerOptions.Create()` for default JSON setup.
- Removed redundant inline JSON configuration in `EnvelopeOptions`.
---
 .../DefaultLambdaJsonSerializerOptions.cs     | 35 +++++++++++++++++++
 .../Options/EnvelopeOptions.cs                | 26 ++------------
 2 files changed, 37 insertions(+), 24 deletions(-)
 create mode 100644 src/AwsLambda.Host.Abstractions/Options/DefaultLambdaJsonSerializerOptions.cs

diff --git a/src/AwsLambda.Host.Abstractions/Options/DefaultLambdaJsonSerializerOptions.cs b/src/AwsLambda.Host.Abstractions/Options/DefaultLambdaJsonSerializerOptions.cs
new file mode 100644
index 00000000..c81de4a4
--- /dev/null
+++ b/src/AwsLambda.Host.Abstractions/Options/DefaultLambdaJsonSerializerOptions.cs
@@ -0,0 +1,35 @@
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using Amazon.Lambda.Serialization.SystemTextJson;
+using Amazon.Lambda.Serialization.SystemTextJson.Converters;
+
+namespace AwsLambda.Host.Options;
+
+/// <summary>Provides the default JSON serializer options used by AWS Lambda.</summary>
+public static class DefaultLambdaJsonSerializerOptions
+{
+    /// <summary>
+    ///     Creates a <see cref="JsonSerializerOptions" /> instance that matches the defaults used by
+    ///     <see cref="DefaultLambdaJsonSerializer" />.
+    /// </summary>
+    /// <remarks>
+    ///     <para>Configures null-value ignoring, case-insensitive property names, and the AWS naming policy.</para>
+    ///     <para>Adds the AWS-provided converters for dates, memory streams, constant classes, and byte arrays.</para>
+    /// </remarks>
+    /// <returns>Configured JSON serializer options suitable for AWS Lambda payloads.</returns>
+    public static JsonSerializerOptions Create()
+    {
+        var options = new JsonSerializerOptions
+        {
+            DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
+            PropertyNameCaseInsensitive = true,
+            PropertyNamingPolicy = new AwsNamingPolicy(),
+        };
+        options.Converters.Add(new DateTimeConverter());
+        options.Converters.Add(new MemoryStreamConverter());
+        options.Converters.Add(new ConstantClassConverter());
+        options.Converters.Add(new ByteArrayConverter());
+
+        return options;
+    }
+}
diff --git a/src/AwsLambda.Host.Abstractions/Options/EnvelopeOptions.cs b/src/AwsLambda.Host.Abstractions/Options/EnvelopeOptions.cs
index 8dc0cc2e..56799568 100644
--- a/src/AwsLambda.Host.Abstractions/Options/EnvelopeOptions.cs
+++ b/src/AwsLambda.Host.Abstractions/Options/EnvelopeOptions.cs
@@ -1,8 +1,6 @@
 using System.Text.Json;
-using System.Text.Json.Serialization;
 using System.Xml;
 using Amazon.Lambda.Serialization.SystemTextJson;
-using Amazon.Lambda.Serialization.SystemTextJson.Converters;
 using AwsLambda.Host.Envelopes;
 
 namespace AwsLambda.Host.Options;
@@ -57,28 +55,8 @@ public class EnvelopeOptions
     ///         not been explicitly configured.
     ///     </para>
     /// </remarks>
-    public JsonSerializerOptions LambdaDefaultJsonOptions
-    {
-        get
-        {
-            if (field is null)
-            {
-                field = new JsonSerializerOptions
-                {
-                    DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
-                    PropertyNameCaseInsensitive = true,
-                    PropertyNamingPolicy = new AwsNamingPolicy(),
-                };
-                field.Converters.Add(new DateTimeConverter());
-                field.Converters.Add(new MemoryStreamConverter());
-                field.Converters.Add(new ConstantClassConverter());
-                field.Converters.Add(new ByteArrayConverter());
-            }
-
-            return field;
-        }
-        set;
-    }
+    public JsonSerializerOptions LambdaDefaultJsonOptions { get; set; } =
+        DefaultLambdaJsonSerializerOptions.Create();
 
     /// <summary>Gets or sets the XML reader settings used when deserializing Lambda event payloads.</summary>
     /// <remarks>

From d78cd62a057ad6aa502542fc558dd2ae9826ea75 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Mon, 8 Dec 2025 20:02:06 -0500
Subject: [PATCH 10/29] refactor(testing): simplify InvocationProcessor and
 update related tests

- Removed unused parameters from `InvocationProcessor` constructor.
- Updated `LambdaApplicationFactory` to use `LambdaServerV2` instead of `LambdaTestServer`.
- Commented out unused test methods in `LambdaHostTest`.
- Modified `Program.cs` to comment out unneeded initialization logic.
- Enhanced `LambdaServerV2` with improved request processing methods.
- Added private methods to handle internal server logic in `LambdaServerV2`.
---
 .../Lambda/Program.cs                         |   8 +-
 .../Tests/LambdaHostTest.cs                   | 255 +++++++++---------
 .../Internal/InvocationProcessor.cs           |   9 +-
 .../LambdaApplicationFactory.cs               |  23 +-
 src/AwsLambda.Host.Testing/LambdaServerV2.cs  | 145 +++++++++-
 5 files changed, 277 insertions(+), 163 deletions(-)

diff --git a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
index b5fa7ee5..4ca272d2 100644
--- a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
@@ -13,10 +13,10 @@
 // Build the Lambda application
 var lambda = builder.Build();
 
-lambda.OnInit(() =>
-{
-    return false;
-});
+// lambda.OnInit(() =>
+// {
+//     return false;
+// });
 
 // Map your handler - the event is automatically injected
 lambda.MapHandler(([Event] string name) => $"Hello {name}!");
diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 7b400362..b7878d52 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -1,7 +1,5 @@
-using AwsLambda.Host.Options;
 using AwsLambda.Host.Testing;
 using JetBrains.Annotations;
-using Microsoft.Extensions.DependencyInjection;
 using Xunit;
 
 namespace Lambda.Host.Example.HelloWorld;
@@ -16,134 +14,131 @@ public async Task LambdaHost_CanStartWithoutError()
 
         var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
-        if (!setup.InitSuccess)
-            throw new Exception("Failed to initialize Lambda host");
-
-        var x = 0;
-
-        var response = await factory.Server.InvokeAsync<string, string>(
-            "Jonas",
-            TestContext.Current.CancellationToken
-        );
-        Assert.True(response.WasSuccess);
-        Assert.NotNull(response);
-        Assert.Equal("Hello Jonas!", response.Response);
-    }
-
-    [Fact]
-    public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
-    {
-        await using var factory = new LambdaApplicationFactory<Program>().WithHostBuilder(builder =>
-        {
-            builder.ConfigureServices(
-                (_, services) =>
-                {
-                    services.Configure<LambdaHostOptions>(options =>
-                    {
-                        options.BootstrapOptions.RuntimeApiEndpoint = "http://localhost:3002";
-                    });
-                }
-            );
-        });
-
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-
-        var response = await factory.Server.InvokeAsync<string, string>(
-            "Jonas",
-            TestContext.Current.CancellationToken
-        );
-        Assert.True(response.WasSuccess);
-        Assert.NotNull(response);
-        Assert.Equal("Hello Jonas!", response.Response);
-    }
-
-    [Fact]
-    public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
-    {
-        await using var factory = new LambdaApplicationFactory<Program>();
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-
-        // Launch 5 concurrent invocations
-        var tasks = Enumerable
-            .Range(1, 5)
-            .Select(i =>
-                factory.Server.InvokeAsync<string, string>(
-                    $"User{i}",
-                    TestContext.Current.CancellationToken
-                )
-            )
-            .ToArray();
-
-        var responses = await Task.WhenAll(tasks);
-
-        // All should complete successfully
-        Assert.All(responses, r => Assert.True(r.WasSuccess));
-        Assert.Equal("Hello User1!", responses[0].Response);
-        Assert.Equal("Hello User2!", responses[1].Response);
-        Assert.Equal("Hello User3!", responses[2].Response);
-        Assert.Equal("Hello User4!", responses[3].Response);
-        Assert.Equal("Hello User5!", responses[4].Response);
-    }
-
-    [Fact]
-    public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
-    {
-        await using var factory = new LambdaApplicationFactory<Program>();
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-
-        var response = await factory.Server.InvokeAsync<string, int>(
-            123,
-            TestContext.Current.CancellationToken
-        );
-
-        Assert.False(response.WasSuccess);
-        Assert.NotNull(response.Error);
-        Assert.Contains("Json", response.Error!.ErrorType, StringComparison.OrdinalIgnoreCase);
-    }
-
-    [Fact]
-    public async Task InvokeAsync_WithPreCanceledToken_CancelsInvocation()
-    {
-        await using var factory = new LambdaApplicationFactory<Program>();
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-
-        using var cts = new CancellationTokenSource();
-        await cts.CancelAsync();
-
-        await Assert.ThrowsAsync<TaskCanceledException>(() =>
-            factory.Server.InvokeAsync<string, string>("Jonas", cts.Token)
-        );
+        // var response = await factory.Server.InvokeAsync<string, string>(
+        //     "Jonas",
+        //     TestContext.Current.CancellationToken
+        // );
+        // Assert.True(response.WasSuccess);
+        // Assert.NotNull(response);
+        // Assert.Equal("Hello Jonas!", response.Response);
     }
 
-    [Fact]
-    public async Task InvokeAsync_WithZeroTimeout_CancelsInvocation() =>
-        await Assert.ThrowsAsync<TaskCanceledException>(async () =>
-        {
-            await using var factory = new LambdaApplicationFactory<Program>();
-            await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-
-            var options = new LambdaClientOptions();
-            options.InvocationHeaderOptions.ClientWaitTimeout = TimeSpan.Zero;
-
-            await factory.Server.InvokeAsync<string, string>(
-                "Jonas",
-                options,
-                TestContext.Current.CancellationToken
-            );
-        });
-
-    [Fact]
-    public async Task StartAsync_WithFailingInit_ReturnsInitError()
-    {
-        // This test verifies that when OnInit returns false (as configured in Program.cs),
-        // the runtime posts to /runtime/init/error and StartAsync returns InitResponse with error
-        await using var factory = new LambdaApplicationFactory<Program>();
-
-        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
-        var initResponse = await factory.Server.StartAsync(cts.Token);
-
-        Assert.False(initResponse.InitSuccess);
-        Assert.NotNull(initResponse.Error);
-        Assert.Equal(ServerState.Stopped, factory.Server.State);
-    }
+    // [Fact]
+    // public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
+    // {
+    //     await using var factory = new LambdaApplicationFactory<Program>().WithHostBuilder(builder
+    // =>
+    //     {
+    //         builder.ConfigureServices(
+    //             (_, services) =>
+    //             {
+    //                 services.Configure<LambdaHostOptions>(options =>
+    //                 {
+    //                     options.BootstrapOptions.RuntimeApiEndpoint = "http://localhost:3002";
+    //                 });
+    //             }
+    //         );
+    //     });
+    //
+    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+    //
+    //     var response = await factory.Server.InvokeAsync<string, string>(
+    //         "Jonas",
+    //         TestContext.Current.CancellationToken
+    //     );
+    //     Assert.True(response.WasSuccess);
+    //     Assert.NotNull(response);
+    //     Assert.Equal("Hello Jonas!", response.Response);
+    // }
+    //
+    // [Fact]
+    // public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
+    // {
+    //     await using var factory = new LambdaApplicationFactory<Program>();
+    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+    //
+    //     // Launch 5 concurrent invocations
+    //     var tasks = Enumerable
+    //         .Range(1, 5)
+    //         .Select(i =>
+    //             factory.Server.InvokeAsync<string, string>(
+    //                 $"User{i}",
+    //                 TestContext.Current.CancellationToken
+    //             )
+    //         )
+    //         .ToArray();
+    //
+    //     var responses = await Task.WhenAll(tasks);
+    //
+    //     // All should complete successfully
+    //     Assert.All(responses, r => Assert.True(r.WasSuccess));
+    //     Assert.Equal("Hello User1!", responses[0].Response);
+    //     Assert.Equal("Hello User2!", responses[1].Response);
+    //     Assert.Equal("Hello User3!", responses[2].Response);
+    //     Assert.Equal("Hello User4!", responses[3].Response);
+    //     Assert.Equal("Hello User5!", responses[4].Response);
+    // }
+    //
+    // [Fact]
+    // public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
+    // {
+    //     await using var factory = new LambdaApplicationFactory<Program>();
+    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+    //
+    //     var response = await factory.Server.InvokeAsync<string, int>(
+    //         123,
+    //         TestContext.Current.CancellationToken
+    //     );
+    //
+    //     Assert.False(response.WasSuccess);
+    //     Assert.NotNull(response.Error);
+    //     Assert.Contains("Json", response.Error!.ErrorType, StringComparison.OrdinalIgnoreCase);
+    // }
+    //
+    // [Fact]
+    // public async Task InvokeAsync_WithPreCanceledToken_CancelsInvocation()
+    // {
+    //     await using var factory = new LambdaApplicationFactory<Program>();
+    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+    //
+    //     using var cts = new CancellationTokenSource();
+    //     await cts.CancelAsync();
+    //
+    //     await Assert.ThrowsAsync<TaskCanceledException>(() =>
+    //         factory.Server.InvokeAsync<string, string>("Jonas", cts.Token)
+    //     );
+    // }
+    //
+    // [Fact]
+    // public async Task InvokeAsync_WithZeroTimeout_CancelsInvocation() =>
+    //     await Assert.ThrowsAsync<TaskCanceledException>(async () =>
+    //     {
+    //         await using var factory = new LambdaApplicationFactory<Program>();
+    //         await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+    //
+    //         var options = new LambdaClientOptions();
+    //         options.InvocationHeaderOptions.ClientWaitTimeout = TimeSpan.Zero;
+    //
+    //         await factory.Server.InvokeAsync<string, string>(
+    //             "Jonas",
+    //             options,
+    //             TestContext.Current.CancellationToken
+    //         );
+    //     });
+    //
+    // [Fact]
+    // public async Task StartAsync_WithFailingInit_ReturnsInitError()
+    // {
+    //     // This test verifies that when OnInit returns false (as configured in Program.cs),
+    //     // the runtime posts to /runtime/init/error and StartAsync returns InitResponse with
+    // error
+    //     await using var factory = new LambdaApplicationFactory<Program>();
+    //
+    //     using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+    //     var initResponse = await factory.Server.StartAsync(cts.Token);
+    //
+    //     Assert.False(initResponse.InitSuccess);
+    //     Assert.NotNull(initResponse.Error);
+    //     Assert.Equal(ServerState.Stopped, factory.Server.State);
+    // }
 }
diff --git a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
index 7e4f503d..e63faed9 100644
--- a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
+++ b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
@@ -27,10 +27,7 @@ internal class InvocationProcessor : IAsyncDisposable
     private ProcessorState _state;
     private Task? _processingTask;
 
-    internal InvocationProcessor(
-        JsonSerializerOptions? jsonSerializerOptions = null,
-        ILambdaRuntimeRouteManager? routeManager = null
-    )
+    internal InvocationProcessor()
     {
         _transactionChannel = Channel.CreateBounded<LambdaHttpTransaction>(
             new BoundedChannelOptions(TransactionChannelCapacity)
@@ -42,8 +39,8 @@ internal InvocationProcessor(
         );
         _pendingInvocationIds = new ConcurrentQueue<string>();
         _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
-        _routeManager = routeManager ?? new LambdaRuntimeRouteManager();
-        _jsonSerializerOptions = jsonSerializerOptions ?? new JsonSerializerOptions();
+        _routeManager = new LambdaRuntimeRouteManager();
+        _jsonSerializerOptions = new JsonSerializerOptions();
         _shutdownCts = new CancellationTokenSource();
         _initCompletionTcs = new TaskCompletionSource<InitResponse>(
             TaskCreationOptions.RunContinuationsAsynchronously
diff --git a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
index c44d9e59..82d0acaa 100644
--- a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
+++ b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
@@ -30,13 +30,13 @@ public class LambdaApplicationFactory<TEntryPoint> : IDisposable, IAsyncDisposab
     private bool _disposed;
     private bool _disposedAsync;
     private IHost? _host;
-    private LambdaTestServer? _server;
+    private LambdaServerV2? _server;
 
     /// <summary>
     /// <para>
     /// Creates an instance of <see cref="LambdaApplicationFactory{TEntryPoint}"/>. This factory can be used to
-    /// create a <see cref="LambdaTestServer"/> instance using the MVC application defined by <typeparamref name="TEntryPoint"/>
-    /// and one or more <see cref="HttpClient"/> instances used to send <see cref="HttpRequestMessage"/> to the <see cref="LambdaTestServer"/>.
+    /// create a <see cref="LambdaServerV2"/> instance using the MVC application defined by <typeparamref name="TEntryPoint"/>
+    /// and one or more <see cref="HttpClient"/> instances used to send <see cref="HttpRequestMessage"/> to the <see cref="LambdaServerV2"/>.
     /// The <see cref="LambdaApplicationFactory{TEntryPoint}"/> will find the entry point class of <typeparamref name="TEntryPoint"/>
     /// assembly and initialize the application by calling <c>IHostBuilder CreateWebHostBuilder(string [] args)</c>
     /// on <typeparamref name="TEntryPoint"/>.
@@ -71,9 +71,9 @@ public class LambdaApplicationFactory<TEntryPoint> : IDisposable, IAsyncDisposab
         _derivedFactories.AsReadOnly();
 
     /// <summary>
-    /// Gets the <see cref="LambdaTestServer"/> created by this <see cref="LambdaApplicationFactory{TEntryPoint}"/>.
+    /// Gets the <see cref="LambdaServerV2"/> created by this <see cref="LambdaApplicationFactory{TEntryPoint}"/>.
     /// </summary>
-    public LambdaTestServer Server
+    public LambdaServerV2 Server
     {
         get
         {
@@ -231,11 +231,10 @@ private void ConfigureHostBuilder(
         _host = CreateHost(hostBuilder);
 
         // Create the public server with the built (but not started) host
-        _server = new LambdaTestServer(
-            _host,
-            processor,
-            entryPointCompletion: entryPointCompletion
-        );
+        // _server = new LambdaServerV2(_host, processor, entryPointCompletion:
+        // entryPointCompletion);
+
+        _server = new LambdaServerV2(_host);
     }
 
     private void SetContentRoot(IHostBuilder builder)
@@ -402,7 +401,7 @@ private static void EnsureDepsFile()
     }
 
     /// <summary>
-    /// Creates a <see cref="IHostBuilder"/> used to set up <see cref="LambdaTestServer"/>.
+    /// Creates a <see cref="IHostBuilder"/> used to set up <see cref="LambdaServerV2"/>.
     /// </summary>
     /// <remarks>
     /// The default implementation of this method looks for a <c>public static IHostBuilder CreateHostBuilder(string[] args)</c>
@@ -428,7 +427,7 @@ private static void EnsureDepsFile()
     /// <param name="builder">The <see cref="IHostBuilder"/> used to create the host.</param>
     /// <returns>The <see cref="IHost"/> with the bootstrapped application.</returns>
     protected virtual IHost CreateHost(IHostBuilder builder) =>
-        // Build the host but DON'T start it - LambdaTestServer.StartAsync() will start it
+        // Build the host but DON'T start it - LambdaServerV2.StartAsync() will start it
         builder.Build();
 
     /// <summary>
diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index 651c953b..d9aa1ecc 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -1,55 +1,178 @@
+using System.Collections.Concurrent;
+using System.Text.Json;
+using System.Threading.Channels;
+using Microsoft.AspNetCore.Routing;
+using Microsoft.Extensions.Hosting;
+
 namespace AwsLambda.Host.Testing;
 
 public class LambdaServerV2 : IAsyncDisposable
 {
-    private bool _started;
-    private bool _stopped;
+    private readonly TaskCompletionSource<InitResponse> _initCompletionTcs;
+    private readonly SemaphoreSlim _invocationAddedSignal;
+    private readonly JsonSerializerOptions _jsonSerializerOptions;
+    private readonly ConcurrentQueue<string> _pendingInvocationIds;
+    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
+    private readonly ILambdaRuntimeRouteManager _routeManager;
+    private readonly CancellationTokenSource _shutdownCts;
+    private readonly Channel<LambdaHttpTransaction> _transactionChannel;
+    private readonly IHost _host;
+
+    private Task? _processingTask;
+    private ServerState _state;
+
+    internal LambdaServerV2(IHost host, CancellationToken shutdownToken = default)
+    {
+        ArgumentNullException.ThrowIfNull(host);
+
+        _host = host;
+
+        _transactionChannel = Channel.CreateUnbounded<LambdaHttpTransaction>(
+            new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
+        );
+        _pendingInvocationIds = new ConcurrentQueue<string>();
+        _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
+        _routeManager = new LambdaRuntimeRouteManager();
+        _jsonSerializerOptions = new JsonSerializerOptions();
+        _shutdownCts = CancellationTokenSource.CreateLinkedTokenSource(shutdownToken);
+        _initCompletionTcs = new TaskCompletionSource<InitResponse>(
+            TaskCreationOptions.RunContinuationsAsynchronously
+        );
+        _invocationAddedSignal = new SemaphoreSlim(0);
+        _state = ServerState.Created;
+    }
 
     public async ValueTask DisposeAsync()
     {
-        if (!_stopped)
-            await StopAsync();
+        await StopAsync();
+
+        _transactionChannel.Writer.TryComplete();
 
-        // TODO release managed resources here
+        _state = ServerState.Disposed;
     }
 
-    public static LambdaServerV2 Create() => new();
+    public IServiceProvider Services => _host.Services;
 
     //      ┌──────────────────────────────────────────────────────────┐
     //      │                        Public API                        │
     //      └──────────────────────────────────────────────────────────┘
 
-    public async Task StartAsync(CancellationToken cancellationToken = default) => _started = true;
+    public async Task<InitResponse> StartAsync(CancellationToken cancellationToken = default)
+    {
+        if (_state != ServerState.Created)
+            throw new InvalidOperationException("Server is already started.");
+
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+
+        _state = ServerState.Starting;
+
+        // Start the host
+        await _host.StartAsync(cts.Token);
+
+        // Start background processing
+        _processingTask = Task.Run(ProcessTransactionsAsync, cts.Token);
+
+        _state = ServerState.Running;
+
+        return default;
+    }
 
     public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
         TEvent invokeEvent,
         CancellationToken cancellationToken = default
     )
     {
-        if (!_started)
+        if (_state != ServerState.Running)
             throw new InvalidOperationException("Server is not started.");
 
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+
         return default;
     }
 
     public async Task StopAsync(CancellationToken cancellationToken = default)
     {
-        if (!_started)
+        if (_state != ServerState.Running)
             return;
 
-        _stopped = true;
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+
+        _state = ServerState.Stopped;
     }
 
     //      ┌──────────────────────────────────────────────────────────┐
     //      │                  Internal Server Logic                   │
     //      └──────────────────────────────────────────────────────────┘
+
+    private async Task ProcessTransactionsAsync()
+    {
+        await foreach (
+            var transaction in _transactionChannel.Reader.ReadAllAsync(_shutdownCts.Token)
+        )
+            try
+            {
+                if (
+                    !_routeManager.TryMatch(
+                        transaction.Request,
+                        out var requestType,
+                        out var routeValues
+                    )
+                )
+                    throw new InvalidOperationException(
+                        $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
+                    );
+
+                switch (requestType!.Value)
+                {
+                    case RequestType.GetNextInvocation:
+                        await HandleGetNextInvocationAsync(transaction);
+                        break;
+
+                    case RequestType.PostResponse:
+                        await HandlePostResponseAsync(transaction, routeValues!);
+                        break;
+
+                    case RequestType.PostError:
+                        await HandlePostErrorAsync(transaction, routeValues!);
+                        break;
+
+                    case RequestType.PostInitError:
+                        await HandlePostInitErrorAsync(transaction);
+                        break;
+
+                    default:
+                        throw new InvalidOperationException(
+                            $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
+                        );
+                }
+            }
+            catch (Exception ex)
+            {
+                // Fail the transaction and continue processing
+                transaction.Fail(ex);
+            }
+    }
+
+    private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction) { }
+
+    private async Task HandlePostResponseAsync(
+        LambdaHttpTransaction transaction,
+        RouteValueDictionary routeValues
+    ) { }
+
+    private async Task HandlePostErrorAsync(
+        LambdaHttpTransaction transaction,
+        RouteValueDictionary routeValues
+    ) { }
+
+    private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction) { }
 }
 
 public static class Temp
 {
     public static async Task Run()
     {
-        await using var server = LambdaServerV2.Create();
+        await using var server = new LambdaServerV2(null!);
         await server.StartAsync();
         var result = await server.InvokeAsync<string, string>("Jonas", CancellationToken.None);
         await server.StopAsync();

From e0b1fc7ea9440757ecf4b465b3d064e60e47b8a7 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 13:50:10 -0500
Subject: [PATCH 11/29] refactor(testing): update InvocationProcessor and
 enhance LambdaServerV2 initialization

- Replaced `InitSuccess` with `InitStatus` in `InitResponse` for more detailed server initialization status.
- Enhanced `InvocationProcessor` state transitions and updated related handling to accommodate `InitStatus`.
- Refactored `LambdaServerV2` to improve lifecycle methods and handle entry point completion gracefully.
- Reinstated and updated `LambdaHostTest` method to verify exception handling on bad configuration.
- Added `InitStatus` enum to standardize server initialization state representation.
---
 .../Tests/LambdaHostTest.cs                   |  59 +++++----
 .../Internal/InvocationProcessor.cs           |   6 +-
 .../LambdaApplicationFactory.cs               |   7 +-
 src/AwsLambda.Host.Testing/LambdaServerV2.cs  | 121 +++++++++++++++---
 .../LambdaTestServer.cs                       |   6 +-
 .../Models/InitResponse.cs                    |  20 ++-
 6 files changed, 164 insertions(+), 55 deletions(-)

diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index b7878d52..5bdcabdb 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -1,5 +1,7 @@
+using AwsLambda.Host.Options;
 using AwsLambda.Host.Testing;
 using JetBrains.Annotations;
+using Microsoft.Extensions.DependencyInjection;
 using Xunit;
 
 namespace Lambda.Host.Example.HelloWorld;
@@ -14,6 +16,8 @@ public async Task LambdaHost_CanStartWithoutError()
 
         var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
+        await Task.Delay(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+
         // var response = await factory.Server.InvokeAsync<string, string>(
         //     "Jonas",
         //     TestContext.Current.CancellationToken
@@ -23,34 +27,33 @@ public async Task LambdaHost_CanStartWithoutError()
         // Assert.Equal("Hello Jonas!", response.Response);
     }
 
-    // [Fact]
-    // public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
-    // {
-    //     await using var factory = new LambdaApplicationFactory<Program>().WithHostBuilder(builder
-    // =>
-    //     {
-    //         builder.ConfigureServices(
-    //             (_, services) =>
-    //             {
-    //                 services.Configure<LambdaHostOptions>(options =>
-    //                 {
-    //                     options.BootstrapOptions.RuntimeApiEndpoint = "http://localhost:3002";
-    //                 });
-    //             }
-    //         );
-    //     });
-    //
-    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-    //
-    //     var response = await factory.Server.InvokeAsync<string, string>(
-    //         "Jonas",
-    //         TestContext.Current.CancellationToken
-    //     );
-    //     Assert.True(response.WasSuccess);
-    //     Assert.NotNull(response);
-    //     Assert.Equal("Hello Jonas!", response.Response);
-    // }
-    //
+    [Fact]
+    public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
+    {
+        await using var factory = new LambdaApplicationFactory<Program>().WithHostBuilder(builder =>
+        {
+            builder.ConfigureServices(
+                (_, services) =>
+                {
+                    services.Configure<LambdaHostOptions>(options =>
+                    {
+                        options.BootstrapOptions.RuntimeApiEndpoint = "http://localhost:3002";
+                    });
+                }
+            );
+        });
+
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        var response = await factory.Server.InvokeAsync<string, string>(
+            "Jonas",
+            TestContext.Current.CancellationToken
+        );
+        Assert.True(response.WasSuccess);
+        Assert.NotNull(response);
+        Assert.Equal("Hello Jonas!", response.Response);
+    }
+
     // [Fact]
     // public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
     // {
diff --git a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
index e63faed9..af2b829d 100644
--- a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
+++ b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
@@ -210,7 +210,9 @@ private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transactio
         if (_state == ProcessorState.Initializing)
         {
             _state = ProcessorState.Running;
-            _initCompletionTcs.TrySetResult(new InitResponse { InitSuccess = true });
+            _initCompletionTcs.TrySetResult(
+                new InitResponse { InitStatus = InitStatus.InitCompleted }
+            );
         }
 
         // Loop until we find work or shutdown
@@ -398,7 +400,7 @@ out var values
         {
             _state = ProcessorState.Stopped;
             _initCompletionTcs.TrySetResult(
-                new InitResponse { InitSuccess = false, Error = errorResponse }
+                new InitResponse { InitStatus = InitStatus.InitError, Error = errorResponse }
             );
             // Wake up any waiting /next requests by releasing the semaphore
             // They will check state and fail appropriately
diff --git a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
index 82d0acaa..2fa533ab 100644
--- a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
+++ b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
@@ -213,12 +213,12 @@ private void ConfigureHostBuilder(
         SetContentRoot(hostBuilder);
         _configuration(hostBuilder);
 
-        var processor = new InvocationProcessor();
+        _server = new LambdaServerV2(entryPointCompletion);
 
         // set Lambda Bootstrap Http Client
         hostBuilder.ConfigureServices(services =>
         {
-            services.AddLambdaBootstrapHttpClient(new HttpClient(processor.CreateTestingHandler()));
+            services.AddLambdaBootstrapHttpClient(new HttpClient(_server.CreateHandler()));
 
             services.PostConfigure<LambdaHostOptions>(options =>
             {
@@ -229,12 +229,11 @@ private void ConfigureHostBuilder(
 
         // Build the host but DON'T start it - server will start it
         _host = CreateHost(hostBuilder);
+        _server.SetHost(_host);
 
         // Create the public server with the built (but not started) host
         // _server = new LambdaServerV2(_host, processor, entryPointCompletion:
         // entryPointCompletion);
-
-        _server = new LambdaServerV2(_host);
     }
 
     private void SetContentRoot(IHostBuilder builder)
diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index d9aa1ecc..3710f0fd 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -1,4 +1,5 @@
 using System.Collections.Concurrent;
+using System.Net.Http.Json;
 using System.Text.Json;
 using System.Threading.Channels;
 using Microsoft.AspNetCore.Routing;
@@ -16,16 +17,21 @@ public class LambdaServerV2 : IAsyncDisposable
     private readonly ILambdaRuntimeRouteManager _routeManager;
     private readonly CancellationTokenSource _shutdownCts;
     private readonly Channel<LambdaHttpTransaction> _transactionChannel;
-    private readonly IHost _host;
+    private readonly Task<Exception?> _entryPointCompletion;
+
+    private IHost? _host;
 
     private Task? _processingTask;
     private ServerState _state;
 
-    internal LambdaServerV2(IHost host, CancellationToken shutdownToken = default)
+    internal LambdaServerV2(
+        Task<Exception?>? entryPointCompletion,
+        CancellationToken shutdownToken = default
+    )
     {
-        ArgumentNullException.ThrowIfNull(host);
+        ArgumentNullException.ThrowIfNull(entryPointCompletion);
 
-        _host = host;
+        _entryPointCompletion = entryPointCompletion;
 
         _transactionChannel = Channel.CreateUnbounded<LambdaHttpTransaction>(
             new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
@@ -42,6 +48,12 @@ internal LambdaServerV2(IHost host, CancellationToken shutdownToken = default)
         _state = ServerState.Created;
     }
 
+    internal void SetHost(IHost host)
+    {
+        ArgumentNullException.ThrowIfNull(host);
+        _host = host;
+    }
+
     public async ValueTask DisposeAsync()
     {
         await StopAsync();
@@ -53,6 +65,9 @@ public async ValueTask DisposeAsync()
 
     public IServiceProvider Services => _host.Services;
 
+    internal HttpMessageHandler CreateHandler() =>
+        new LambdaTestingHttpHandler(_transactionChannel);
+
     //      ┌──────────────────────────────────────────────────────────┐
     //      │                        Public API                        │
     //      └──────────────────────────────────────────────────────────┘
@@ -62,6 +77,9 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
         if (_state != ServerState.Created)
             throw new InvalidOperationException("Server is already started.");
 
+        if (_host is null)
+            throw new InvalidOperationException("Host is not set.");
+
         using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
         _state = ServerState.Starting;
@@ -72,11 +90,53 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
         // Start background processing
         _processingTask = Task.Run(ProcessTransactionsAsync, cts.Token);
 
-        _state = ServerState.Running;
+        var exceptions = await WhenAny(
+            _processingTask,
+            _entryPointCompletion,
+            _initCompletionTcs.Task
+        );
 
-        return default;
+        if (exceptions.Length > 0)
+            throw exceptions.Length > 0
+                ? new AggregateException(
+                    "Multiple exceptions encountered while running StartAsync",
+                    exceptions
+                )
+                : exceptions[0];
+
+        if (_entryPointCompletion.IsCompleted)
+            return new InitResponse { InitStatus = InitStatus.HostExited };
+
+        if (_initCompletionTcs.Task.IsCompleted)
+        {
+            _state =
+                _initCompletionTcs.Task.Result.InitStatus == InitStatus.InitCompleted
+                    ? ServerState.Running
+                    : ServerState.Stopped;
+            return _initCompletionTcs.Task.Result;
+        }
+
+        throw new InvalidOperationException(
+            "Server initialization failed with neither an error nor completion."
+        );
     }
 
+    private static async Task<Exception[]> WhenAny(params Task[] tasks)
+    {
+        await Task.WhenAny(tasks);
+        return ExtractExceptions(tasks);
+    }
+
+    private static Exception[] ExtractExceptions(Task[] tasks) =>
+        tasks
+            .Where(t => t is { IsFaulted: true, Exception: not null })
+            .Select(e =>
+                e.Exception!.InnerExceptions.Count > 1
+                    ? e.Exception!
+                    : e.Exception.InnerExceptions[0]
+            )
+            .ToArray();
+
     public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
         TEvent invokeEvent,
         CancellationToken cancellationToken = default
@@ -149,11 +209,18 @@ out var routeValues
             catch (Exception ex)
             {
                 // Fail the transaction and continue processing
-                transaction.Fail(ex);
+                // transaction.Fail(ex);
+                throw;
             }
     }
 
-    private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction) { }
+    private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)
+    {
+        if (_state == ServerState.Starting)
+            _initCompletionTcs.SetResult(
+                new InitResponse { InitStatus = InitStatus.InitCompleted }
+            );
+    }
 
     private async Task HandlePostResponseAsync(
         LambdaHttpTransaction transaction,
@@ -165,16 +232,34 @@ private async Task HandlePostErrorAsync(
         RouteValueDictionary routeValues
     ) { }
 
-    private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction) { }
-}
-
-public static class Temp
-{
-    public static async Task Run()
+    private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction)
     {
-        await using var server = new LambdaServerV2(null!);
-        await server.StartAsync();
-        var result = await server.InvokeAsync<string, string>("Jonas", CancellationToken.None);
-        await server.StopAsync();
+        if (_state == ServerState.Starting)
+            _initCompletionTcs.SetResult(
+                new InitResponse
+                {
+                    Error = await (
+                        transaction.Request.Content?.ReadFromJsonAsync<ErrorResponse>(
+                            _jsonSerializerOptions
+                        ) ?? Task.FromResult<ErrorResponse?>(null)
+                    ),
+                    InitStatus = InitStatus.InitError,
+                }
+            );
+
+        throw new InvalidOperationException(
+            "Server is already started and as such an initialization error cannot be reported."
+        );
     }
 }
+
+// public static class Temp
+// {
+//     public static async Task Run()
+//     {
+//         await using var server = new LambdaServerV2();
+//         await server.StartAsync();
+//         var result = await server.InvokeAsync<string, string>("Jonas", CancellationToken.None);
+//         await server.StopAsync();
+//     }
+// }
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServer.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
index 2d8048ad..d5d0fdfd 100644
--- a/src/AwsLambda.Host.Testing/LambdaTestServer.cs
+++ b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
@@ -95,8 +95,10 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
             var initResponse = await initTask;
 
             // If init failed, auto-stop the server
-            if (!initResponse.InitSuccess)
-                await StopAsync(CancellationToken.None);
+            // if (initResponse != InitStatus.InitCompleted)
+            //     await StopAsync(CancellationToken.None);
+            // if (!initResponse.InitSuccess)
+            //     await StopAsync(CancellationToken.None);
 
             return initResponse;
         }
diff --git a/src/AwsLambda.Host.Testing/Models/InitResponse.cs b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
index d861159a..c3034302 100644
--- a/src/AwsLambda.Host.Testing/Models/InitResponse.cs
+++ b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
@@ -3,5 +3,23 @@ namespace AwsLambda.Host.Testing;
 public class InitResponse
 {
     public ErrorResponse? Error { get; internal set; }
-    public bool InitSuccess { get; internal set; }
+    public InitStatus InitStatus { get; internal set; }
+}
+
+public enum InitStatus
+{
+    /// <summary>
+    /// Initialization of the Lambda completed successfully.
+    /// </summary>
+    InitCompleted,
+
+    /// <summary>
+    /// Initialization of the Lambda failed, and the Lambda returned an error.
+    /// </summary>
+    InitError,
+
+    /// <summary>
+    /// Initialization of the Lambda failed, and the Host process exited.
+    /// </summary>
+    HostExited,
 }

From 793abd97de7d56ab403ad9b8bb97742170347ac2 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 17:15:22 -0500
Subject: [PATCH 12/29] test(testing): reinstate and update LambdaHost tests
 for functionality verification

- Re-enabled and updated `LambdaHost_ProcessesConcurrentInvocationsInFifoOrder` test for concurrency validation.
- Reinstated logic in `LambdaHostTest` to verify response handling and assertion correctness.
- Updated comments in `Program.cs` and adjusted initialization-related code.
- Added XML documentation to `InitResponse` and `InitStatus` for clarity and better maintainability.
- Refactored `LambdaServerV2` with internal method enhancements and exception processing optimizations.
---
 .../Lambda/Program.cs                         |  5 +-
 .../Tests/LambdaHostTest.cs                   | 72 +++++++++----------
 src/AwsLambda.Host.Testing/LambdaServerV2.cs  | 14 +++-
 .../Models/InitResponse.cs                    | 17 ++++-
 4 files changed, 68 insertions(+), 40 deletions(-)

diff --git a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
index 4ca272d2..137329b6 100644
--- a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
@@ -13,9 +13,12 @@
 // Build the Lambda application
 var lambda = builder.Build();
 
+// throw new Exception("Init failed");
+
 // lambda.OnInit(() =>
 // {
-//     return false;
+//     // throw new Exception("Init failed");
+//     // return false;
 // });
 
 // Map your handler - the event is automatically injected
diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 5bdcabdb..f965b071 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -16,15 +16,15 @@ public async Task LambdaHost_CanStartWithoutError()
 
         var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
-        await Task.Delay(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+        // await Task.Delay(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
 
-        // var response = await factory.Server.InvokeAsync<string, string>(
-        //     "Jonas",
-        //     TestContext.Current.CancellationToken
-        // );
-        // Assert.True(response.WasSuccess);
-        // Assert.NotNull(response);
-        // Assert.Equal("Hello Jonas!", response.Response);
+        var response = await factory.Server.InvokeAsync<string, string>(
+            "Jonas",
+            TestContext.Current.CancellationToken
+        );
+        Assert.True(response.WasSuccess);
+        Assert.NotNull(response);
+        Assert.Equal("Hello Jonas!", response.Response);
     }
 
     [Fact]
@@ -54,34 +54,34 @@ public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
         Assert.Equal("Hello Jonas!", response.Response);
     }
 
-    // [Fact]
-    // public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
-    // {
-    //     await using var factory = new LambdaApplicationFactory<Program>();
-    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-    //
-    //     // Launch 5 concurrent invocations
-    //     var tasks = Enumerable
-    //         .Range(1, 5)
-    //         .Select(i =>
-    //             factory.Server.InvokeAsync<string, string>(
-    //                 $"User{i}",
-    //                 TestContext.Current.CancellationToken
-    //             )
-    //         )
-    //         .ToArray();
-    //
-    //     var responses = await Task.WhenAll(tasks);
-    //
-    //     // All should complete successfully
-    //     Assert.All(responses, r => Assert.True(r.WasSuccess));
-    //     Assert.Equal("Hello User1!", responses[0].Response);
-    //     Assert.Equal("Hello User2!", responses[1].Response);
-    //     Assert.Equal("Hello User3!", responses[2].Response);
-    //     Assert.Equal("Hello User4!", responses[3].Response);
-    //     Assert.Equal("Hello User5!", responses[4].Response);
-    // }
-    //
+    [Fact]
+    public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
+    {
+        await using var factory = new LambdaApplicationFactory<Program>();
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        // Launch 5 concurrent invocations
+        var tasks = Enumerable
+            .Range(1, 5)
+            .Select(i =>
+                factory.Server.InvokeAsync<string, string>(
+                    $"User{i}",
+                    TestContext.Current.CancellationToken
+                )
+            )
+            .ToArray();
+
+        var responses = await Task.WhenAll(tasks);
+
+        // All should complete successfully
+        Assert.All(responses, r => Assert.True(r.WasSuccess));
+        Assert.Equal("Hello User1!", responses[0].Response);
+        Assert.Equal("Hello User2!", responses[1].Response);
+        Assert.Equal("Hello User3!", responses[2].Response);
+        Assert.Equal("Hello User4!", responses[3].Response);
+        Assert.Equal("Hello User5!", responses[4].Response);
+    }
+
     // [Fact]
     // public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
     // {
diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index 3710f0fd..7532c9b5 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -113,6 +113,7 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
                 _initCompletionTcs.Task.Result.InitStatus == InitStatus.InitCompleted
                     ? ServerState.Running
                     : ServerState.Stopped;
+
             return _initCompletionTcs.Task.Result;
         }
 
@@ -127,12 +128,18 @@ private static async Task<Exception[]> WhenAny(params Task[] tasks)
         return ExtractExceptions(tasks);
     }
 
+    private static async Task<Exception[]> WhenAll(params Task[] tasks)
+    {
+        await Task.WhenAll(tasks);
+        return ExtractExceptions(tasks);
+    }
+
     private static Exception[] ExtractExceptions(Task[] tasks) =>
         tasks
             .Where(t => t is { IsFaulted: true, Exception: not null })
             .Select(e =>
                 e.Exception!.InnerExceptions.Count > 1
-                    ? e.Exception!
+                    ? e.Exception
                     : e.Exception.InnerExceptions[0]
             )
             .ToArray();
@@ -147,6 +154,8 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
 
         using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
+        var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
+
         return default;
     }
 
@@ -217,9 +226,12 @@ out var routeValues
     private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)
     {
         if (_state == ServerState.Starting)
+        {
             _initCompletionTcs.SetResult(
                 new InitResponse { InitStatus = InitStatus.InitCompleted }
             );
+            return;
+        }
     }
 
     private async Task HandlePostResponseAsync(
diff --git a/src/AwsLambda.Host.Testing/Models/InitResponse.cs b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
index c3034302..c15b754f 100644
--- a/src/AwsLambda.Host.Testing/Models/InitResponse.cs
+++ b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
@@ -1,11 +1,24 @@
 namespace AwsLambda.Host.Testing;
 
+/// <summary>
+/// Represents the result of a Lambda function initialization attempt.
+/// </summary>
 public class InitResponse
 {
-    public ErrorResponse? Error { get; internal set; }
-    public InitStatus InitStatus { get; internal set; }
+    /// <summary>
+    /// Gets the error information if initialization failed, or null if initialization succeeded.
+    /// </summary>
+    public ErrorResponse? Error { get; internal init; }
+
+    /// <summary>
+    /// Gets the status of the initialization attempt.
+    /// </summary>
+    public InitStatus InitStatus { get; internal init; }
 }
 
+/// <summary>
+/// An enumeration of possible statuses for Lambda initialization.
+/// </summary>
 public enum InitStatus
 {
     /// <summary>

From 15817623df5536eb4e37574209330f2da9250323 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 19:01:26 -0500
Subject: [PATCH 13/29] feat(testing): enhance LambdaServerV2 request handling
 and header management

- Introduced `_clientOptions` field to centralize Lambda client configuration.
- Added `GetRequestId` method for unique request ID generation with padding.
- Implemented `CreateEventResponse` to generate event responses with headers and JSON content.
- Replaced `ConcurrentDictionary` with `Channel` for `PendingInvocation` management.
- Enhanced `InvokeAsync` to process and return invocation responses with error handling.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 98 +++++++++++++++++++-
 1 file changed, 94 insertions(+), 4 deletions(-)

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index 7532c9b5..467c5fe5 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -1,5 +1,7 @@
 using System.Collections.Concurrent;
+using System.Net;
 using System.Net.Http.Json;
+using System.Text;
 using System.Text.Json;
 using System.Threading.Channels;
 using Microsoft.AspNetCore.Routing;
@@ -13,14 +15,16 @@ public class LambdaServerV2 : IAsyncDisposable
     private readonly SemaphoreSlim _invocationAddedSignal;
     private readonly JsonSerializerOptions _jsonSerializerOptions;
     private readonly ConcurrentQueue<string> _pendingInvocationIds;
-    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
+    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations_old;
     private readonly ILambdaRuntimeRouteManager _routeManager;
     private readonly CancellationTokenSource _shutdownCts;
     private readonly Channel<LambdaHttpTransaction> _transactionChannel;
     private readonly Task<Exception?> _entryPointCompletion;
+    private readonly LambdaClientOptions _clientOptions;
+    private readonly Channel<PendingInvocation> _pendingInvocations;
 
     private IHost? _host;
-
+    private int _requestCounter;
     private Task? _processingTask;
     private ServerState _state;
 
@@ -36,8 +40,11 @@ internal LambdaServerV2(
         _transactionChannel = Channel.CreateUnbounded<LambdaHttpTransaction>(
             new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
         );
+        _pendingInvocations = Channel.CreateUnbounded<PendingInvocation>(
+            new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
+        );
         _pendingInvocationIds = new ConcurrentQueue<string>();
-        _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
+        // _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
         _routeManager = new LambdaRuntimeRouteManager();
         _jsonSerializerOptions = new JsonSerializerOptions();
         _shutdownCts = CancellationTokenSource.CreateLinkedTokenSource(shutdownToken);
@@ -46,6 +53,7 @@ internal LambdaServerV2(
         );
         _invocationAddedSignal = new SemaphoreSlim(0);
         _state = ServerState.Created;
+        _clientOptions = new LambdaClientOptions();
     }
 
     internal void SetHost(IHost host)
@@ -154,9 +162,48 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
 
         using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
+        // Generate unique request ID
+        var requestId = GetRequestId();
+
+        // Create the event response with Lambda headers
+        var eventResponse = CreateEventResponse(invokeEvent, requestId);
+        var deadlineUtc = DateTimeOffset.UtcNow.Add(
+            _clientOptions.InvocationHeaderOptions.ClientWaitTimeout
+        );
+
+        cts.CancelAfter(_clientOptions.InvocationHeaderOptions.ClientWaitTimeout);
+
         var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
 
-        return default;
+        _pendingInvocations.Writer.TryWrite(pending);
+
+        var completion = await pending.ResponseTcs.Task.WaitAsync(cts.Token);
+
+        var responseMessage = completion.Request;
+        var wasSuccess = completion.RequestType == RequestType.PostResponse;
+
+        var invocationResponse = new InvocationResponse<TResponse>
+        {
+            WasSuccess = wasSuccess,
+            Response = wasSuccess
+                ? await (
+                    responseMessage.Content?.ReadFromJsonAsync<TResponse>(
+                        _jsonSerializerOptions,
+                        cts.Token
+                    ) ?? Task.FromResult<TResponse?>(default)
+                )
+                : default,
+            Error = !wasSuccess
+                ? await (
+                    responseMessage.Content?.ReadFromJsonAsync<ErrorResponse>(
+                        _jsonSerializerOptions,
+                        cts.Token
+                    ) ?? Task.FromResult<ErrorResponse?>(null)
+                )
+                : null,
+        };
+
+        return invocationResponse;
     }
 
     public async Task StopAsync(CancellationToken cancellationToken = default)
@@ -263,6 +310,49 @@ private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction)
             "Server is already started and as such an initialization error cannot be reported."
         );
     }
+
+    private HttpResponseMessage CreateEventResponse<TEvent>(TEvent invokeEvent, string requestId)
+    {
+        var response = new HttpResponseMessage(HttpStatusCode.OK)
+        {
+            Content = new StringContent(
+                JsonSerializer.Serialize(invokeEvent, _jsonSerializerOptions),
+                Encoding.UTF8,
+                "application/json"
+            ),
+            Version = Version.Parse("1.1"),
+        };
+
+        // Add standard HTTP headers
+        response.Headers.Date = new DateTimeOffset(_clientOptions.InvocationHeaderOptions.Date);
+        response.Headers.TransferEncodingChunked = _clientOptions
+            .InvocationHeaderOptions
+            .TransferEncodingChunked;
+
+        // Add custom Lambda runtime headers
+        var deadlineMs = DateTimeOffset
+            .UtcNow.Add(_clientOptions.InvocationHeaderOptions.FunctionTimeout)
+            .ToUnixTimeMilliseconds();
+        response.Headers.Add("Lambda-Runtime-Deadline-Ms", deadlineMs.ToString());
+        response.Headers.Add("Lambda-Runtime-Aws-Request-Id", requestId);
+        response.Headers.Add(
+            "Lambda-Runtime-Trace-Id",
+            _clientOptions.InvocationHeaderOptions.TraceId
+        );
+        response.Headers.Add(
+            "Lambda-Runtime-Invoked-Function-Arn",
+            _clientOptions.InvocationHeaderOptions.FunctionArn
+        );
+
+        // Add any additional custom headers
+        foreach (var header in _clientOptions.InvocationHeaderOptions.AdditionalHeaders)
+            response.Headers.Add(header.Key, header.Value);
+
+        return response;
+    }
+
+    private string GetRequestId() =>
+        Interlocked.Increment(ref _requestCounter).ToString().PadLeft(12, '0');
 }
 
 // public static class Temp

From 5f398bdaf1f90813fcb4391fdd24bb71bb3b4b1c Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 20:07:40 -0500
Subject: [PATCH 14/29] feat(testing): refactor LambdaServerV2 and update
 LambdaHostTest assertions

- Simplified `LambdaServerV2` by replacing `ConcurrentQueue` with `Channel` for invocation management.
- Refactored variable initialization and lifecycle handling for clarity and reuse.
- Updated `LambdaHostTest` to use Fluent Assertions for improved readability and maintainability.
- Introduced a private method, `CreateCompletionAsync`, for cloning and processing HTTP requests.
- Added safeguards against duplicate request IDs in `LambdaServerV2`.
---
 .../Tests/LambdaHostTest.cs                   |  10 +-
 src/AwsLambda.Host.Testing/LambdaServerV2.cs  | 169 ++++++++++++------
 2 files changed, 117 insertions(+), 62 deletions(-)

diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index f965b071..8b2e3fea 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -1,3 +1,4 @@
+using AwesomeAssertions;
 using AwsLambda.Host.Options;
 using AwsLambda.Host.Testing;
 using JetBrains.Annotations;
@@ -16,15 +17,16 @@ public async Task LambdaHost_CanStartWithoutError()
 
         var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
-        // await Task.Delay(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
+        setup.InitStatus.Should().Be(InitStatus.InitCompleted);
 
         var response = await factory.Server.InvokeAsync<string, string>(
             "Jonas",
             TestContext.Current.CancellationToken
         );
-        Assert.True(response.WasSuccess);
-        Assert.NotNull(response);
-        Assert.Equal("Hello Jonas!", response.Response);
+
+        response.WasSuccess.Should().BeTrue();
+        response.Should().NotBeNull();
+        response.Response.Should().Be("Hello Jonas!");
     }
 
     [Fact]
diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index 467c5fe5..af8a00b7 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -11,21 +11,20 @@ namespace AwsLambda.Host.Testing;
 
 public class LambdaServerV2 : IAsyncDisposable
 {
+    private readonly LambdaClientOptions _clientOptions;
+    private readonly Task<Exception?> _entryPointCompletion;
     private readonly TaskCompletionSource<InitResponse> _initCompletionTcs;
     private readonly SemaphoreSlim _invocationAddedSignal;
     private readonly JsonSerializerOptions _jsonSerializerOptions;
-    private readonly ConcurrentQueue<string> _pendingInvocationIds;
-    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations_old;
+    private readonly Channel<string> _pendingInvocationIds;
+    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
     private readonly ILambdaRuntimeRouteManager _routeManager;
     private readonly CancellationTokenSource _shutdownCts;
     private readonly Channel<LambdaHttpTransaction> _transactionChannel;
-    private readonly Task<Exception?> _entryPointCompletion;
-    private readonly LambdaClientOptions _clientOptions;
-    private readonly Channel<PendingInvocation> _pendingInvocations;
 
     private IHost? _host;
-    private int _requestCounter;
     private Task? _processingTask;
+    private int _requestCounter;
     private ServerState _state;
 
     internal LambdaServerV2(
@@ -40,11 +39,10 @@ internal LambdaServerV2(
         _transactionChannel = Channel.CreateUnbounded<LambdaHttpTransaction>(
             new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
         );
-        _pendingInvocations = Channel.CreateUnbounded<PendingInvocation>(
+        _pendingInvocationIds = Channel.CreateUnbounded<string>(
             new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
         );
-        _pendingInvocationIds = new ConcurrentQueue<string>();
-        // _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
+        _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
         _routeManager = new LambdaRuntimeRouteManager();
         _jsonSerializerOptions = new JsonSerializerOptions();
         _shutdownCts = CancellationTokenSource.CreateLinkedTokenSource(shutdownToken);
@@ -56,11 +54,7 @@ internal LambdaServerV2(
         _clientOptions = new LambdaClientOptions();
     }
 
-    internal void SetHost(IHost host)
-    {
-        ArgumentNullException.ThrowIfNull(host);
-        _host = host;
-    }
+    public IServiceProvider Services => _host.Services;
 
     public async ValueTask DisposeAsync()
     {
@@ -71,7 +65,11 @@ public async ValueTask DisposeAsync()
         _state = ServerState.Disposed;
     }
 
-    public IServiceProvider Services => _host.Services;
+    internal void SetHost(IHost host)
+    {
+        ArgumentNullException.ThrowIfNull(host);
+        _host = host;
+    }
 
     internal HttpMessageHandler CreateHandler() =>
         new LambdaTestingHttpHandler(_transactionChannel);
@@ -175,7 +173,9 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
 
         var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
 
-        _pendingInvocations.Writer.TryWrite(pending);
+        if (!_pendingInvocations.TryAdd(requestId, pending))
+            throw new InvalidOperationException($"Duplicate request ID: {requestId}");
+        _pendingInvocationIds.Writer.TryWrite(requestId);
 
         var completion = await pending.ResponseTcs.Task.WaitAsync(cts.Token);
 
@@ -225,66 +225,88 @@ private async Task ProcessTransactionsAsync()
         await foreach (
             var transaction in _transactionChannel.Reader.ReadAllAsync(_shutdownCts.Token)
         )
-            try
-            {
-                if (
-                    !_routeManager.TryMatch(
-                        transaction.Request,
-                        out var requestType,
-                        out var routeValues
-                    )
+        {
+            if (
+                !_routeManager.TryMatch(
+                    transaction.Request,
+                    out var requestType,
+                    out var routeValues
                 )
-                    throw new InvalidOperationException(
-                        $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
-                    );
+            )
+                throw new InvalidOperationException(
+                    $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
+                );
 
-                switch (requestType!.Value)
-                {
-                    case RequestType.GetNextInvocation:
-                        await HandleGetNextInvocationAsync(transaction);
-                        break;
-
-                    case RequestType.PostResponse:
-                        await HandlePostResponseAsync(transaction, routeValues!);
-                        break;
-
-                    case RequestType.PostError:
-                        await HandlePostErrorAsync(transaction, routeValues!);
-                        break;
-
-                    case RequestType.PostInitError:
-                        await HandlePostInitErrorAsync(transaction);
-                        break;
-
-                    default:
-                        throw new InvalidOperationException(
-                            $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
-                        );
-                }
-            }
-            catch (Exception ex)
+            switch (requestType!.Value)
             {
-                // Fail the transaction and continue processing
-                // transaction.Fail(ex);
-                throw;
+                case RequestType.GetNextInvocation:
+                    await HandleGetNextInvocationAsync(transaction);
+                    break;
+
+                case RequestType.PostResponse:
+                    await HandlePostResponseAsync(transaction, routeValues!);
+                    break;
+
+                case RequestType.PostError:
+                    await HandlePostErrorAsync(transaction, routeValues!);
+                    break;
+
+                case RequestType.PostInitError:
+                    await HandlePostInitErrorAsync(transaction);
+                    break;
+
+                default:
+                    throw new InvalidOperationException(
+                        $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
+                    );
             }
+        }
     }
 
     private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)
     {
         if (_state == ServerState.Starting)
-        {
             _initCompletionTcs.SetResult(
                 new InitResponse { InitStatus = InitStatus.InitCompleted }
             );
-            return;
+
+        if (await _pendingInvocationIds.Reader.WaitToReadAsync(_shutdownCts.Token))
+        {
+            var requestId = await _pendingInvocationIds.Reader.ReadAsync(_shutdownCts.Token);
+            if (!_pendingInvocations.TryGetValue(requestId, out var pendingInvocation))
+                throw new InvalidOperationException($"Missing pending invocation for {requestId}");
+            transaction.ResponseTcs.SetResult(pendingInvocation.EventResponse);
         }
     }
 
     private async Task HandlePostResponseAsync(
         LambdaHttpTransaction transaction,
         RouteValueDictionary routeValues
-    ) { }
+    )
+    {
+        var requestId = routeValues["requestId"]?.ToString();
+        if (requestId is null || !_pendingInvocations.TryGetValue(requestId, out var pending))
+            throw new InvalidOperationException($"Missing pending invocation for {requestId}");
+
+        // Acknowledge to Bootstrap
+        transaction.Respond(
+            new HttpResponseMessage(HttpStatusCode.Accepted)
+            {
+                Content = new StringContent(
+                    """
+                    {"status":"success"}
+                    """,
+                    Encoding.UTF8,
+                    "application/json"
+                ),
+                Version = Version.Parse("1.1"),
+            }
+        );
+
+        pending.ResponseTcs.SetResult(
+            await CreateCompletionAsync(RequestType.PostResponse, transaction.Request)
+        );
+    }
 
     private async Task HandlePostErrorAsync(
         LambdaHttpTransaction transaction,
@@ -353,6 +375,37 @@ private HttpResponseMessage CreateEventResponse<TEvent>(TEvent invokeEvent, stri
 
     private string GetRequestId() =>
         Interlocked.Increment(ref _requestCounter).ToString().PadLeft(12, '0');
+
+    private static async Task<InvocationCompletion> CreateCompletionAsync(
+        RequestType requestType,
+        HttpRequestMessage sourceRequest
+    )
+    {
+        var clonedRequest = new HttpRequestMessage(sourceRequest.Method, sourceRequest.RequestUri)
+        {
+            Version = sourceRequest.Version,
+            VersionPolicy = sourceRequest.VersionPolicy,
+        };
+
+        foreach (var header in sourceRequest.Headers)
+            clonedRequest.Headers.TryAddWithoutValidation(header.Key, header.Value);
+
+        foreach (var option in sourceRequest.Options)
+            clonedRequest.Options.TryAdd(option.Key, option.Value);
+
+        if (sourceRequest.Content != null)
+        {
+            var contentBytes = await sourceRequest.Content.ReadAsByteArrayAsync();
+            var clonedContent = new ByteArrayContent(contentBytes);
+
+            foreach (var header in sourceRequest.Content.Headers)
+                clonedContent.Headers.TryAddWithoutValidation(header.Key, header.Value);
+
+            clonedRequest.Content = clonedContent;
+        }
+
+        return new InvocationCompletion { Request = clonedRequest, RequestType = requestType };
+    }
 }
 
 // public static class Temp

From 4e786ab99fea6741be8b2036146a8ce049d59b56 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 20:15:34 -0500
Subject: [PATCH 15/29] refactor(testing): extract common success response
 creation in LambdaServerV2

- Introduced `CreateSuccessResponse` method to centralize success response generation.
- Updated `HandlePostErrorAsync` and related methods to use `CreateSuccessResponse`.
- Improved code reuse and reduced duplication in response creation logic.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 41 +++++++++++++-------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index af8a00b7..dbc76210 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -289,19 +289,7 @@ RouteValueDictionary routeValues
             throw new InvalidOperationException($"Missing pending invocation for {requestId}");
 
         // Acknowledge to Bootstrap
-        transaction.Respond(
-            new HttpResponseMessage(HttpStatusCode.Accepted)
-            {
-                Content = new StringContent(
-                    """
-                    {"status":"success"}
-                    """,
-                    Encoding.UTF8,
-                    "application/json"
-                ),
-                Version = Version.Parse("1.1"),
-            }
-        );
+        transaction.Respond(CreateSuccessResponse());
 
         pending.ResponseTcs.SetResult(
             await CreateCompletionAsync(RequestType.PostResponse, transaction.Request)
@@ -311,7 +299,19 @@ await CreateCompletionAsync(RequestType.PostResponse, transaction.Request)
     private async Task HandlePostErrorAsync(
         LambdaHttpTransaction transaction,
         RouteValueDictionary routeValues
-    ) { }
+    )
+    {
+        var requestId = routeValues["requestId"]?.ToString();
+        if (requestId is null || !_pendingInvocations.TryGetValue(requestId, out var pending))
+            throw new InvalidOperationException($"Missing pending invocation for {requestId}");
+
+        // Acknowledge to Bootstrap
+        transaction.Respond(CreateSuccessResponse());
+
+        pending.ResponseTcs.SetResult(
+            await CreateCompletionAsync(RequestType.PostError, transaction.Request)
+        );
+    }
 
     private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction)
     {
@@ -406,6 +406,19 @@ HttpRequestMessage sourceRequest
 
         return new InvocationCompletion { Request = clonedRequest, RequestType = requestType };
     }
+
+    private static HttpResponseMessage CreateSuccessResponse() =>
+        new(HttpStatusCode.Accepted)
+        {
+            Content = new StringContent(
+                """
+                {"status":"success"}
+                """,
+                Encoding.UTF8,
+                "application/json"
+            ),
+            Version = Version.Parse("1.1"),
+        };
 }
 
 // public static class Temp

From 91816a6f892cd06925a89798fb822e870f58327a Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 20:15:45 -0500
Subject: [PATCH 16/29] test(lambda-host): add test case for LambdaHost error
 response validation

- Introduced `LambdaHost_HandlerReturnsError` to verify error handling with missing name parameter.
- Asserted error response properties, including error message and success flag.
- Ensured proper server initialization and invocation flow.
---
 .../Tests/LambdaHostTest.cs                   | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 8b2e3fea..b1931373 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -29,6 +29,25 @@ public async Task LambdaHost_CanStartWithoutError()
         response.Response.Should().Be("Hello Jonas!");
     }
 
+    [Fact]
+    public async Task LambdaHost_HandlerReturnsError()
+    {
+        await using var factory = new LambdaApplicationFactory<Program>();
+
+        var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        setup.InitStatus.Should().Be(InitStatus.InitCompleted);
+
+        var response = await factory.Server.InvokeAsync<string, string>(
+            "",
+            TestContext.Current.CancellationToken
+        );
+
+        response.WasSuccess.Should().BeFalse();
+        response.Error.Should().NotBeNull();
+        response.Error?.ErrorMessage.Should().Be("Name is required. (Parameter 'name')");
+    }
+
     [Fact]
     public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
     {

From 138c06055273d5f7f043c064058c82bdd3a0b85c Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 20:16:01 -0500
Subject: [PATCH 17/29] feat(lambda-host): add validation for empty name
 parameter in handler

- Updated handler to throw `ArgumentNullException` if the name parameter is null or whitespace.
- Enhanced input validation to ensure proper error handling during Lambda execution.
---
 .../AwsLambda.Host.Example.Testing/Lambda/Program.cs   | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
index 137329b6..8d51efed 100644
--- a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
@@ -22,7 +22,15 @@
 // });
 
 // Map your handler - the event is automatically injected
-lambda.MapHandler(([Event] string name) => $"Hello {name}!");
+lambda.MapHandler(
+    ([Event] string name) =>
+    {
+        if (string.IsNullOrWhiteSpace(name))
+            throw new ArgumentNullException(nameof(name), "Name is required.");
+
+        return $"Hello {name}!";
+    }
+);
 
 // Run the Lambda
 await lambda.RunAsync();

From fbcaccd7f9b6632a89fe46b3174d934ef56f13a5 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 20:30:57 -0500
Subject: [PATCH 18/29] feat(testing): enhance LambdaServerV2 lifecycle and
 application management

- Added `_applicationLifetime` to manage host application lifecycle.
- Updated `Services` property to use null-forgiving operator on `_host`.
- Integrated `_applicationLifetime.StopApplication()` during server shutdown.
- Ensured proper server stop handling with `_entryPointCompletion` await logic.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index dbc76210..332f4119 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -5,6 +5,7 @@
 using System.Text.Json;
 using System.Threading.Channels;
 using Microsoft.AspNetCore.Routing;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 
 namespace AwsLambda.Host.Testing;
@@ -21,6 +22,7 @@ public class LambdaServerV2 : IAsyncDisposable
     private readonly ILambdaRuntimeRouteManager _routeManager;
     private readonly CancellationTokenSource _shutdownCts;
     private readonly Channel<LambdaHttpTransaction> _transactionChannel;
+    private IHostApplicationLifetime _applicationLifetime;
 
     private IHost? _host;
     private Task? _processingTask;
@@ -54,7 +56,7 @@ internal LambdaServerV2(
         _clientOptions = new LambdaClientOptions();
     }
 
-    public IServiceProvider Services => _host.Services;
+    public IServiceProvider Services => _host!.Services;
 
     public async ValueTask DisposeAsync()
     {
@@ -90,6 +92,8 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
 
         _state = ServerState.Starting;
 
+        _applicationLifetime = _host.Services.GetRequiredService<IHostApplicationLifetime>();
+
         // Start the host
         await _host.StartAsync(cts.Token);
 
@@ -213,6 +217,10 @@ public async Task StopAsync(CancellationToken cancellationToken = default)
 
         using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
+        _applicationLifetime.StopApplication();
+
+        await _entryPointCompletion;
+
         _state = ServerState.Stopped;
     }
 

From 329bcff0dca0796ef3485b34b189927d04eadef9 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Tue, 9 Dec 2025 20:31:07 -0500
Subject: [PATCH 19/29] test(lambda-host): add test for verifying LambdaHost
 shutdown process

- Introduced `LambdaHost_CanBeShutdown` to verify proper server initialization and shutdown handling.
- Asserted initialization status, response properties, and server shutdown behavior.
- Added `OnShutdown` callback in `Program.cs` to log shutdown event.
---
 .../Lambda/Program.cs                         |  5 +++++
 .../Tests/LambdaHostTest.cs                   | 21 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
index 8d51efed..f5e0ed0d 100644
--- a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
@@ -32,6 +32,11 @@
     }
 );
 
+lambda.OnShutdown(() =>
+{
+    Console.WriteLine("Shutdown");
+});
+
 // Run the Lambda
 await lambda.RunAsync();
 
diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index b1931373..42ff17a4 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -48,6 +48,27 @@ public async Task LambdaHost_HandlerReturnsError()
         response.Error?.ErrorMessage.Should().Be("Name is required. (Parameter 'name')");
     }
 
+    [Fact]
+    public async Task LambdaHost_CanBeShutdown()
+    {
+        await using var factory = new LambdaApplicationFactory<Program>();
+
+        var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        setup.InitStatus.Should().Be(InitStatus.InitCompleted);
+
+        var response = await factory.Server.InvokeAsync<string, string>(
+            "Jonas",
+            TestContext.Current.CancellationToken
+        );
+
+        response.WasSuccess.Should().BeTrue();
+        response.Should().NotBeNull();
+        response.Response.Should().Be("Hello Jonas!");
+
+        await factory.Server.StopAsync(TestContext.Current.CancellationToken);
+    }
+
     [Fact]
     public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
     {

From 7358470431d64c0ffe1686f895c457be899a8344 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 10:15:43 -0500
Subject: [PATCH 20/29] feat(testing): enhance LambdaServerV2 lifecycle and
 shutdown process

- Improved member initialization in `LambdaServerV2` for clarity and consistency.
- Added XML documentation to fields and methods for better code maintainability.
- Updated shutdown logic to use `_shutdownCts.CancelAsync()` for proper task cancellation.
- Refactored unused or duplicated logic in constructor.
- Enhanced server state management during shutdown to enforce lifecycle rules.
- Integrated `DefaultLambdaJsonSerializerOptions.Create()` for standardized JSON serialization.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 97 +++++++++++++++-----
 1 file changed, 75 insertions(+), 22 deletions(-)

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index 332f4119..a1b589ee 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -4,6 +4,7 @@
 using System.Text;
 using System.Text.Json;
 using System.Threading.Channels;
+using AwsLambda.Host.Options;
 using Microsoft.AspNetCore.Routing;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
@@ -12,21 +13,81 @@ namespace AwsLambda.Host.Testing;
 
 public class LambdaServerV2 : IAsyncDisposable
 {
+    /// <summary>
+    /// Options used to configure how the server interacts with the Lambda.
+    /// </summary>
     private readonly LambdaClientOptions _clientOptions;
+
+    /// <summary>
+    /// Task that represents the running Host application that has been captured.
+    /// </summary>
     private readonly Task<Exception?> _entryPointCompletion;
-    private readonly TaskCompletionSource<InitResponse> _initCompletionTcs;
-    private readonly SemaphoreSlim _invocationAddedSignal;
+
+    /// <summary>
+    /// TCS used to signal the startup has completed
+    /// </summary>
+    private readonly TaskCompletionSource<InitResponse> _initCompletionTcs = new(
+        TaskCreationOptions.RunContinuationsAsynchronously
+    );
+
+    /// <summary>
+    /// JSON serializer options used to serialize/deserilize Lambda events and responses.
+    /// </summary>
     private readonly JsonSerializerOptions _jsonSerializerOptions;
-    private readonly Channel<string> _pendingInvocationIds;
-    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
-    private readonly ILambdaRuntimeRouteManager _routeManager;
+
+    /// <summary>
+    /// Channel used to queue pending invocations in a FIFO manner.
+    /// </summary>
+    private readonly Channel<string> _pendingInvocationIds = Channel.CreateUnbounded<string>(
+        new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
+    );
+
+    /// <summary>
+    /// Dictionary to track all invocations that have been sent to Lambda.
+    /// </summary>
+    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations = new();
+
+    /// <summary>
+    /// Route manager to determine the route of the incoming request from the Lambda.
+    /// </summary>
+    private readonly ILambdaRuntimeRouteManager _routeManager = new LambdaRuntimeRouteManager();
+
+    /// <summary>
+    /// CTS used to signal shutdown of the server and cancellation of pending tasks.
+    /// </summary>
     private readonly CancellationTokenSource _shutdownCts;
-    private readonly Channel<LambdaHttpTransaction> _transactionChannel;
+
+    /// <summary>
+    /// Channel used to by the Lambda to send events to the server.
+    /// </summary>
+    private readonly Channel<LambdaHttpTransaction> _transactionChannel =
+        Channel.CreateUnbounded<LambdaHttpTransaction>(
+            new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
+        );
+
+    /// <summary>
+    /// Host application lifetime used to signal shutdown to the captioned Host.
+    /// </summary>
     private IHostApplicationLifetime _applicationLifetime;
 
+    /// <summary>
+    /// The captured Host instance.
+    /// </summary>
     private IHost? _host;
+
+    /// <summary>
+    /// Task that is running the background processing loop to handle incoming requests from Lambda.
+    /// </summary>
     private Task? _processingTask;
+
+    /// <summary>
+    /// Counter used to generate unique request IDs.
+    /// </summary>
     private int _requestCounter;
+
+    /// <summary>
+    /// Current state of the server used to enforce lifecycle rules.
+    /// </summary>
     private ServerState _state;
 
     internal LambdaServerV2(
@@ -37,22 +98,10 @@ internal LambdaServerV2(
         ArgumentNullException.ThrowIfNull(entryPointCompletion);
 
         _entryPointCompletion = entryPointCompletion;
-
-        _transactionChannel = Channel.CreateUnbounded<LambdaHttpTransaction>(
-            new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
-        );
-        _pendingInvocationIds = Channel.CreateUnbounded<string>(
-            new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }
-        );
-        _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
-        _routeManager = new LambdaRuntimeRouteManager();
-        _jsonSerializerOptions = new JsonSerializerOptions();
         _shutdownCts = CancellationTokenSource.CreateLinkedTokenSource(shutdownToken);
-        _initCompletionTcs = new TaskCompletionSource<InitResponse>(
-            TaskCreationOptions.RunContinuationsAsynchronously
-        );
-        _invocationAddedSignal = new SemaphoreSlim(0);
         _state = ServerState.Created;
+
+        _jsonSerializerOptions = DefaultLambdaJsonSerializerOptions.Create();
         _clientOptions = new LambdaClientOptions();
     }
 
@@ -64,6 +113,8 @@ public async ValueTask DisposeAsync()
 
         _transactionChannel.Writer.TryComplete();
 
+        await _shutdownCts.CancelAsync();
+
         _state = ServerState.Disposed;
     }
 
@@ -215,11 +266,13 @@ public async Task StopAsync(CancellationToken cancellationToken = default)
         if (_state != ServerState.Running)
             return;
 
-        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        _state = ServerState.Stopping;
+
+        await _shutdownCts.CancelAsync();
 
         _applicationLifetime.StopApplication();
 
-        await _entryPointCompletion;
+        var exceptions = await WhenAll(_entryPointCompletion, _processingTask!);
 
         _state = ServerState.Stopped;
     }

From a35f0002b71aac03f404cfa406238d04584de616 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 10:41:37 -0500
Subject: [PATCH 21/29] feat(testing): add `TaskHelpers` utility for task
 exception handling

- Introduced `WhenAny` and `WhenAll` methods to extract exceptions from completed tasks.
- Added `UnwrapAndThrow` extension for awaiting tasks and throwing aggregate exceptions.
- Improved task exception handling for testing scenarios in `AwsLambda.Host.Testing`.
---
 src/AwsLambda.Host.Testing/TaskHelpers.cs | 36 +++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 src/AwsLambda.Host.Testing/TaskHelpers.cs

diff --git a/src/AwsLambda.Host.Testing/TaskHelpers.cs b/src/AwsLambda.Host.Testing/TaskHelpers.cs
new file mode 100644
index 00000000..8d23fd5c
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/TaskHelpers.cs
@@ -0,0 +1,36 @@
+namespace AwsLambda.Host.Testing;
+
+internal static class TaskHelpers
+{
+    internal static async Task<Exception[]> WhenAny(params Task[] tasks)
+    {
+        await Task.WhenAny(tasks);
+        return ExtractExceptions(tasks);
+    }
+
+    internal static async Task<Exception[]> WhenAll(params Task[] tasks)
+    {
+        await Task.WhenAll(tasks);
+        return ExtractExceptions(tasks);
+    }
+
+    private static Exception[] ExtractExceptions(Task[] tasks) =>
+        tasks
+            .Where(t => t is { IsFaulted: true, Exception: not null })
+            .Select(e =>
+                e.Exception!.InnerExceptions.Count > 1
+                    ? e.Exception
+                    : e.Exception.InnerExceptions[0]
+            )
+            .ToArray();
+
+    extension(Task<Exception[]> exceptionsTask)
+    {
+        internal async Task UnwrapAndThrow(string errorMessage)
+        {
+            var exceptions = await exceptionsTask;
+            if (exceptions.Length > 0)
+                throw new AggregateException(errorMessage, exceptions);
+        }
+    }
+}

From 9312c7a0621f8f510be4a325a4aae51bcf22fc2f Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 10:41:52 -0500
Subject: [PATCH 22/29] feat(testing): add DictionaryExtensions for dictionary
 operations

- Introduced `AddRequired` method to ensure unique key addition and throw on duplicates.
- Added `GetRequired` method to retrieve values by key and enforce key existence.
- Improved error handling for dictionary operations in testing scenarios.
---
 .../DictionaryExtensions.cs                   | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)
 create mode 100644 src/AwsLambda.Host.Testing/DictionaryExtensions.cs

diff --git a/src/AwsLambda.Host.Testing/DictionaryExtensions.cs b/src/AwsLambda.Host.Testing/DictionaryExtensions.cs
new file mode 100644
index 00000000..7a8e5afc
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/DictionaryExtensions.cs
@@ -0,0 +1,19 @@
+namespace AwsLambda.Host.Testing;
+
+internal static class DictionaryExtensions
+{
+    extension<TKey, TValue>(IDictionary<TKey, TValue> dictionary)
+    {
+        internal void AddRequired(TKey key, TValue value)
+        {
+            if (!dictionary.TryAdd(key, value))
+                throw new InvalidOperationException($"Key '{key}' already exists.");
+        }
+
+        internal void GetRequired(TKey? key, out TValue value)
+        {
+            if (key is null || !dictionary.TryGetValue(key, out value!))
+                throw new InvalidOperationException($"Key '{key}' is null or does not exist.");
+        }
+    }
+}

From d8628b1f9aa7324deb3ade734fe599235d5fc8b2 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 10:42:03 -0500
Subject: [PATCH 23/29] feat(testing): replace exception-handling utilities and
 improve invocation handling in LambdaServerV2

- Replaced custom exception-extracting methods with `TaskHelpers` utility methods (`WhenAny`, `WhenAll`).
- Simplified exception handling using `UnwrapAndThrow` for consistent error reporting.
- Streamlined `InvokeAsync` response handling by splitting success and error processing.
- Replaced duplicate `TryAdd` and `TryGetValue` calls with `AddRequired` and `GetRequired`.
- Cleaned up unused methods and redundant logic to enhance maintainability.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 108 ++++++-------------
 1 file changed, 32 insertions(+), 76 deletions(-)

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index a1b589ee..d88c4253 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -151,19 +151,9 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
         // Start background processing
         _processingTask = Task.Run(ProcessTransactionsAsync, cts.Token);
 
-        var exceptions = await WhenAny(
-            _processingTask,
-            _entryPointCompletion,
-            _initCompletionTcs.Task
-        );
-
-        if (exceptions.Length > 0)
-            throw exceptions.Length > 0
-                ? new AggregateException(
-                    "Multiple exceptions encountered while running StartAsync",
-                    exceptions
-                )
-                : exceptions[0];
+        await TaskHelpers
+            .WhenAny(_processingTask, _entryPointCompletion, _initCompletionTcs.Task)
+            .UnwrapAndThrow("Exception(s) encountered while running StartAsync");
 
         if (_entryPointCompletion.IsCompleted)
             return new InitResponse { InitStatus = InitStatus.HostExited };
@@ -183,28 +173,6 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
         );
     }
 
-    private static async Task<Exception[]> WhenAny(params Task[] tasks)
-    {
-        await Task.WhenAny(tasks);
-        return ExtractExceptions(tasks);
-    }
-
-    private static async Task<Exception[]> WhenAll(params Task[] tasks)
-    {
-        await Task.WhenAll(tasks);
-        return ExtractExceptions(tasks);
-    }
-
-    private static Exception[] ExtractExceptions(Task[] tasks) =>
-        tasks
-            .Where(t => t is { IsFaulted: true, Exception: not null })
-            .Select(e =>
-                e.Exception!.InnerExceptions.Count > 1
-                    ? e.Exception
-                    : e.Exception.InnerExceptions[0]
-            )
-            .ToArray();
-
     public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
         TEvent invokeEvent,
         CancellationToken cancellationToken = default
@@ -228,8 +196,8 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
 
         var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
 
-        if (!_pendingInvocations.TryAdd(requestId, pending))
-            throw new InvalidOperationException($"Duplicate request ID: {requestId}");
+        _pendingInvocations.AddRequired(requestId, pending);
+
         _pendingInvocationIds.Writer.TryWrite(requestId);
 
         var completion = await pending.ResponseTcs.Task.WaitAsync(cts.Token);
@@ -237,28 +205,30 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
         var responseMessage = completion.Request;
         var wasSuccess = completion.RequestType == RequestType.PostResponse;
 
-        var invocationResponse = new InvocationResponse<TResponse>
+        var response = wasSuccess
+            ? await (
+                responseMessage.Content?.ReadFromJsonAsync<TResponse>(
+                    _jsonSerializerOptions,
+                    cts.Token
+                ) ?? Task.FromResult<TResponse?>(default)
+            )
+            : default;
+
+        var error = !wasSuccess
+            ? await (
+                responseMessage.Content?.ReadFromJsonAsync<ErrorResponse>(
+                    _jsonSerializerOptions,
+                    cts.Token
+                ) ?? Task.FromResult<ErrorResponse?>(null)
+            )
+            : null;
+
+        return new InvocationResponse<TResponse>
         {
             WasSuccess = wasSuccess,
-            Response = wasSuccess
-                ? await (
-                    responseMessage.Content?.ReadFromJsonAsync<TResponse>(
-                        _jsonSerializerOptions,
-                        cts.Token
-                    ) ?? Task.FromResult<TResponse?>(default)
-                )
-                : default,
-            Error = !wasSuccess
-                ? await (
-                    responseMessage.Content?.ReadFromJsonAsync<ErrorResponse>(
-                        _jsonSerializerOptions,
-                        cts.Token
-                    ) ?? Task.FromResult<ErrorResponse?>(null)
-                )
-                : null,
+            Response = response,
+            Error = error,
         };
-
-        return invocationResponse;
     }
 
     public async Task StopAsync(CancellationToken cancellationToken = default)
@@ -272,7 +242,9 @@ public async Task StopAsync(CancellationToken cancellationToken = default)
 
         _applicationLifetime.StopApplication();
 
-        var exceptions = await WhenAll(_entryPointCompletion, _processingTask!);
+        await TaskHelpers
+            .WhenAll(_entryPointCompletion, _processingTask!)
+            .UnwrapAndThrow("Exception(s) encountered while running StopAsync");
 
         _state = ServerState.Stopped;
     }
@@ -334,8 +306,7 @@ private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transactio
         if (await _pendingInvocationIds.Reader.WaitToReadAsync(_shutdownCts.Token))
         {
             var requestId = await _pendingInvocationIds.Reader.ReadAsync(_shutdownCts.Token);
-            if (!_pendingInvocations.TryGetValue(requestId, out var pendingInvocation))
-                throw new InvalidOperationException($"Missing pending invocation for {requestId}");
+            _pendingInvocations.GetRequired(requestId, out var pendingInvocation);
             transaction.ResponseTcs.SetResult(pendingInvocation.EventResponse);
         }
     }
@@ -345,9 +316,7 @@ private async Task HandlePostResponseAsync(
         RouteValueDictionary routeValues
     )
     {
-        var requestId = routeValues["requestId"]?.ToString();
-        if (requestId is null || !_pendingInvocations.TryGetValue(requestId, out var pending))
-            throw new InvalidOperationException($"Missing pending invocation for {requestId}");
+        _pendingInvocations.GetRequired(routeValues["requestId"]?.ToString(), out var pending);
 
         // Acknowledge to Bootstrap
         transaction.Respond(CreateSuccessResponse());
@@ -362,9 +331,7 @@ private async Task HandlePostErrorAsync(
         RouteValueDictionary routeValues
     )
     {
-        var requestId = routeValues["requestId"]?.ToString();
-        if (requestId is null || !_pendingInvocations.TryGetValue(requestId, out var pending))
-            throw new InvalidOperationException($"Missing pending invocation for {requestId}");
+        _pendingInvocations.GetRequired(routeValues["requestId"]?.ToString(), out var pending);
 
         // Acknowledge to Bootstrap
         transaction.Respond(CreateSuccessResponse());
@@ -481,14 +448,3 @@ private static HttpResponseMessage CreateSuccessResponse() =>
             Version = Version.Parse("1.1"),
         };
 }
-
-// public static class Temp
-// {
-//     public static async Task Run()
-//     {
-//         await using var server = new LambdaServerV2();
-//         await server.StartAsync();
-//         var result = await server.InvokeAsync<string, string>("Jonas", CancellationToken.None);
-//         await server.StopAsync();
-//     }
-// }

From 05e39c90b3bbfd7f2b161a849fd0007568b0119b Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 10:45:05 -0500
Subject: [PATCH 24/29] feat(testing): improve error handling and lifecycle
 management in LambdaServerV2

- Added specific error messages for invalid state transitions during lifecycle operations.
- Enhanced `DisposeAsync` to prevent redundant `StopAsync` calls when the server is not running.
- Improved state validation logic with descriptive exceptions for `StartAsync`, `InvokeAsync`, and `StopAsync`.
- Streamlined shutdown behavior to enforce consistent state management and avoid silent failures.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index d88c4253..f4476a39 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -109,7 +109,8 @@ internal LambdaServerV2(
 
     public async ValueTask DisposeAsync()
     {
-        await StopAsync();
+        if (_state == ServerState.Running)
+            await StopAsync();
 
         _transactionChannel.Writer.TryComplete();
 
@@ -134,7 +135,9 @@ internal HttpMessageHandler CreateHandler() =>
     public async Task<InitResponse> StartAsync(CancellationToken cancellationToken = default)
     {
         if (_state != ServerState.Created)
-            throw new InvalidOperationException("Server is already started.");
+            throw new InvalidOperationException(
+                "Server has already been started and cannot be restarted."
+            );
 
         if (_host is null)
             throw new InvalidOperationException("Host is not set.");
@@ -179,7 +182,9 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
     )
     {
         if (_state != ServerState.Running)
-            throw new InvalidOperationException("Server is not started.");
+            throw new InvalidOperationException(
+                "Server is not Running and as such an event cannot be invoked."
+            );
 
         using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
@@ -234,7 +239,9 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
     public async Task StopAsync(CancellationToken cancellationToken = default)
     {
         if (_state != ServerState.Running)
-            return;
+            throw new InvalidOperationException(
+                "Server is not running and as such cannot be stopped."
+            );
 
         _state = ServerState.Stopping;
 

From b037e6ec85f0c7795a082f78fee15d36999a68c9 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 11:05:17 -0500
Subject: [PATCH 25/29] feat(testing): enhance cancellation handling and error
 safety in LambdaServerV2

- Added `_shutdownCts.Token` to `CancellationTokenSource.CreateLinkedTokenSource` for consistency.
- Improved `ProcessTransactionsAsync` with try-catch to handle task cancellations gracefully.
- Ensured proper cancellation behavior in `StopAsync` with `.WaitAsync(cancellationToken)`.
---
 src/AwsLambda.Host.Testing/LambdaServerV2.cs | 84 ++++++++++++--------
 1 file changed, 49 insertions(+), 35 deletions(-)

diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index f4476a39..72fce135 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -142,7 +142,10 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
         if (_host is null)
             throw new InvalidOperationException("Host is not set.");
 
-        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(
+            cancellationToken,
+            _shutdownCts.Token
+        );
 
         _state = ServerState.Starting;
 
@@ -186,7 +189,10 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
                 "Server is not Running and as such an event cannot be invoked."
             );
 
-        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(
+            cancellationToken,
+            _shutdownCts.Token
+        );
 
         // Generate unique request ID
         var requestId = GetRequestId();
@@ -251,7 +257,8 @@ public async Task StopAsync(CancellationToken cancellationToken = default)
 
         await TaskHelpers
             .WhenAll(_entryPointCompletion, _processingTask!)
-            .UnwrapAndThrow("Exception(s) encountered while running StopAsync");
+            .UnwrapAndThrow("Exception(s) encountered while running StopAsync")
+            .WaitAsync(cancellationToken);
 
         _state = ServerState.Stopped;
     }
@@ -262,45 +269,52 @@ await TaskHelpers
 
     private async Task ProcessTransactionsAsync()
     {
-        await foreach (
-            var transaction in _transactionChannel.Reader.ReadAllAsync(_shutdownCts.Token)
-        )
+        try
         {
-            if (
-                !_routeManager.TryMatch(
-                    transaction.Request,
-                    out var requestType,
-                    out var routeValues
-                )
+            await foreach (
+                var transaction in _transactionChannel.Reader.ReadAllAsync(_shutdownCts.Token)
             )
-                throw new InvalidOperationException(
-                    $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
-                );
-
-            switch (requestType!.Value)
             {
-                case RequestType.GetNextInvocation:
-                    await HandleGetNextInvocationAsync(transaction);
-                    break;
-
-                case RequestType.PostResponse:
-                    await HandlePostResponseAsync(transaction, routeValues!);
-                    break;
-
-                case RequestType.PostError:
-                    await HandlePostErrorAsync(transaction, routeValues!);
-                    break;
-
-                case RequestType.PostInitError:
-                    await HandlePostInitErrorAsync(transaction);
-                    break;
-
-                default:
+                if (
+                    !_routeManager.TryMatch(
+                        transaction.Request,
+                        out var requestType,
+                        out var routeValues
+                    )
+                )
                     throw new InvalidOperationException(
-                        $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
+                        $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
                     );
+
+                switch (requestType!.Value)
+                {
+                    case RequestType.GetNextInvocation:
+                        await HandleGetNextInvocationAsync(transaction);
+                        break;
+
+                    case RequestType.PostResponse:
+                        await HandlePostResponseAsync(transaction, routeValues!);
+                        break;
+
+                    case RequestType.PostError:
+                        await HandlePostErrorAsync(transaction, routeValues!);
+                        break;
+
+                    case RequestType.PostInitError:
+                        await HandlePostInitErrorAsync(transaction);
+                        break;
+
+                    default:
+                        throw new InvalidOperationException(
+                            $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
+                        );
+                }
             }
         }
+        catch (OperationCanceledException) when (_shutdownCts.IsCancellationRequested)
+        {
+            // Expected when task is canceled
+        }
     }
 
     private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)

From 185a0fc6882086c247b53f332a48099fbc94f98f Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 11:31:17 -0500
Subject: [PATCH 26/29] feat(testing): remove unused Lambda server testing
 infrastructure

- Deleted `ILambdaRuntimeRouteManager`, `InvocationProcessor`, and `_ProcessorState`.
- Removed `LambdaTestServer` class and its related models and methods.
- Cleared redundant internal classes like `LambdaBootstrapRequest` from `AwsLambda.Host.Testing`.
- Cleaned up unnecessary internal abstractions to streamline testing code and reduce maintenance burden.
---
 .../Lambda/Program.cs                         |   3 +-
 .../Tests/LambdaHostTest.cs                   |  84 +--
 .../ILambdaRuntimeRouteManager.cs             |  12 -
 .../Internal/InvocationProcessor.cs           | 494 ------------------
 .../Internal/ProcessorState.cs                |  35 --
 .../LambdaRuntimeRouteManager.cs              |   2 +-
 src/AwsLambda.Host.Testing/LambdaServerV2.cs  |   2 +-
 .../LambdaTestServer.cs                       | 316 -----------
 .../Models/InitResponse.cs                    |  21 -
 .../Models/InitStatus.cs                      |  22 +
 .../Models/LambdaBootstrapRequest.cs          |  24 -
 .../{ => Models}/ServerState.cs               |   0
 src/AwsLambda.Host.Testing/runtime-api.yaml   | 250 ---------
 13 files changed, 69 insertions(+), 1196 deletions(-)
 delete mode 100644 src/AwsLambda.Host.Testing/ILambdaRuntimeRouteManager.cs
 delete mode 100644 src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
 delete mode 100644 src/AwsLambda.Host.Testing/Internal/ProcessorState.cs
 delete mode 100644 src/AwsLambda.Host.Testing/LambdaTestServer.cs
 create mode 100644 src/AwsLambda.Host.Testing/Models/InitStatus.cs
 delete mode 100644 src/AwsLambda.Host.Testing/Models/LambdaBootstrapRequest.cs
 rename src/AwsLambda.Host.Testing/{ => Models}/ServerState.cs (100%)
 delete mode 100644 src/AwsLambda.Host.Testing/runtime-api.yaml

diff --git a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
index f5e0ed0d..59f49a23 100644
--- a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
@@ -23,8 +23,9 @@
 
 // Map your handler - the event is automatically injected
 lambda.MapHandler(
-    ([Event] string name) =>
+    async ([Event] string name, CancellationToken cancellationToken) =>
     {
+        // await Task.Delay(TimeSpan.FromSeconds(60), cancellationToken);
         if (string.IsNullOrWhiteSpace(name))
             throw new ArgumentNullException(nameof(name), "Name is required.");
 
diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 42ff17a4..5d81f73a 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -70,7 +70,7 @@ public async Task LambdaHost_CanBeShutdown()
     }
 
     [Fact]
-    public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
+    public async Task LambdaHost_ServerInternalExceptions_AreCaughtAndReturnedAsError()
     {
         await using var factory = new LambdaApplicationFactory<Program>().WithHostBuilder(builder =>
         {
@@ -85,15 +85,17 @@ public async Task LambdaHost_CrashesWithBadConfiguration_ThrowsException()
             );
         });
 
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+        var act = async () =>
+            await factory.Server.StartAsync(TestContext.Current.CancellationToken);
 
-        var response = await factory.Server.InvokeAsync<string, string>(
-            "Jonas",
-            TestContext.Current.CancellationToken
-        );
-        Assert.True(response.WasSuccess);
-        Assert.NotNull(response);
-        Assert.Equal("Hello Jonas!", response.Response);
+        (await act.Should().ThrowAsync<AggregateException>())
+            .And.InnerExceptions.Should()
+            .ContainSingle(ex =>
+                ex is InvalidOperationException
+                && ex.Message.Contains(
+                    "Unexpected request: GET http://http//localhost:3002/2018-06-01/runtime/invocation/next"
+                )
+            );
     }
 
     [Fact]
@@ -124,36 +126,36 @@ public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
         Assert.Equal("Hello User5!", responses[4].Response);
     }
 
-    // [Fact]
-    // public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
-    // {
-    //     await using var factory = new LambdaApplicationFactory<Program>();
-    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-    //
-    //     var response = await factory.Server.InvokeAsync<string, int>(
-    //         123,
-    //         TestContext.Current.CancellationToken
-    //     );
-    //
-    //     Assert.False(response.WasSuccess);
-    //     Assert.NotNull(response.Error);
-    //     Assert.Contains("Json", response.Error!.ErrorType, StringComparison.OrdinalIgnoreCase);
-    // }
-    //
-    // [Fact]
-    // public async Task InvokeAsync_WithPreCanceledToken_CancelsInvocation()
-    // {
-    //     await using var factory = new LambdaApplicationFactory<Program>();
-    //     await factory.Server.StartAsync(TestContext.Current.CancellationToken);
-    //
-    //     using var cts = new CancellationTokenSource();
-    //     await cts.CancelAsync();
-    //
-    //     await Assert.ThrowsAsync<TaskCanceledException>(() =>
-    //         factory.Server.InvokeAsync<string, string>("Jonas", cts.Token)
-    //     );
-    // }
-    //
+    [Fact]
+    public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
+    {
+        await using var factory = new LambdaApplicationFactory<Program>();
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        var response = await factory.Server.InvokeAsync<string, int>(
+            123,
+            TestContext.Current.CancellationToken
+        );
+
+        Assert.False(response.WasSuccess);
+        Assert.NotNull(response.Error);
+        Assert.Contains("Json", response.Error!.ErrorType, StringComparison.OrdinalIgnoreCase);
+    }
+
+    [Fact]
+    public async Task InvokeAsync_WithPreCanceledToken_CancelsInvocation()
+    {
+        await using var factory = new LambdaApplicationFactory<Program>();
+        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+
+        using var cts = new CancellationTokenSource();
+        await cts.CancelAsync();
+
+        await Assert.ThrowsAsync<TaskCanceledException>(() =>
+            factory.Server.InvokeAsync<string, string>("Jonas", cts.Token)
+        );
+    }
+
     // [Fact]
     // public async Task InvokeAsync_WithZeroTimeout_CancelsInvocation() =>
     //     await Assert.ThrowsAsync<TaskCanceledException>(async () =>
@@ -170,13 +172,13 @@ public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
     //             TestContext.Current.CancellationToken
     //         );
     //     });
-    //
+
     // [Fact]
     // public async Task StartAsync_WithFailingInit_ReturnsInitError()
     // {
     //     // This test verifies that when OnInit returns false (as configured in Program.cs),
     //     // the runtime posts to /runtime/init/error and StartAsync returns InitResponse with
-    // error
+    //     //error
     //     await using var factory = new LambdaApplicationFactory<Program>();
     //
     //     using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
diff --git a/src/AwsLambda.Host.Testing/ILambdaRuntimeRouteManager.cs b/src/AwsLambda.Host.Testing/ILambdaRuntimeRouteManager.cs
deleted file mode 100644
index 041586b6..00000000
--- a/src/AwsLambda.Host.Testing/ILambdaRuntimeRouteManager.cs
+++ /dev/null
@@ -1,12 +0,0 @@
-using Microsoft.AspNetCore.Routing;
-
-namespace AwsLambda.Host.Testing;
-
-internal interface ILambdaRuntimeRouteManager
-{
-    bool TryMatch(
-        HttpRequestMessage request,
-        out RequestType? routeType,
-        out RouteValueDictionary? values
-    );
-}
diff --git a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs b/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
deleted file mode 100644
index af2b829d..00000000
--- a/src/AwsLambda.Host.Testing/Internal/InvocationProcessor.cs
+++ /dev/null
@@ -1,494 +0,0 @@
-using System.Collections.Concurrent;
-using System.Net;
-using System.Net.Http.Json;
-using System.Text;
-using System.Text.Json;
-using System.Threading.Channels;
-using Microsoft.AspNetCore.Routing;
-
-namespace AwsLambda.Host.Testing;
-
-/// <summary>
-/// Internal processor that handles HTTP transactions from Lambda Bootstrap.
-/// Routes requests, queues invocations, and manages request-response correlation.
-/// </summary>
-internal class InvocationProcessor : IAsyncDisposable
-{
-    private const int TransactionChannelCapacity = 1024;
-
-    private readonly JsonSerializerOptions _jsonSerializerOptions;
-    private readonly ConcurrentQueue<string> _pendingInvocationIds;
-    private readonly ConcurrentDictionary<string, PendingInvocation> _pendingInvocations;
-    private readonly ILambdaRuntimeRouteManager _routeManager;
-    private readonly CancellationTokenSource _shutdownCts;
-    private readonly Channel<LambdaHttpTransaction> _transactionChannel;
-    private readonly TaskCompletionSource<InitResponse> _initCompletionTcs;
-    private readonly SemaphoreSlim _invocationAddedSignal;
-    private ProcessorState _state;
-    private Task? _processingTask;
-
-    internal InvocationProcessor()
-    {
-        _transactionChannel = Channel.CreateBounded<LambdaHttpTransaction>(
-            new BoundedChannelOptions(TransactionChannelCapacity)
-            {
-                SingleReader = true,
-                SingleWriter = false,
-                FullMode = BoundedChannelFullMode.Wait,
-            }
-        );
-        _pendingInvocationIds = new ConcurrentQueue<string>();
-        _pendingInvocations = new ConcurrentDictionary<string, PendingInvocation>();
-        _routeManager = new LambdaRuntimeRouteManager();
-        _jsonSerializerOptions = new JsonSerializerOptions();
-        _shutdownCts = new CancellationTokenSource();
-        _initCompletionTcs = new TaskCompletionSource<InitResponse>(
-            TaskCreationOptions.RunContinuationsAsynchronously
-        );
-        _invocationAddedSignal = new SemaphoreSlim(0);
-        _state = ProcessorState.Created;
-    }
-
-    public async ValueTask DisposeAsync()
-    {
-        // Transition to stopping state
-        if (_state != ProcessorState.Stopped)
-            _state = ProcessorState.Stopping;
-
-        await _shutdownCts.CancelAsync();
-
-        // Complete init if still pending (server shutting down before init)
-        if (_state == ProcessorState.Initializing || _state == ProcessorState.Stopping)
-            _initCompletionTcs.TrySetCanceled(_shutdownCts.Token);
-
-        _transactionChannel.Writer.TryComplete();
-
-        // Fail any in-flight transactions that haven't been processed yet
-        var disposedException = new OperationCanceledException("InvocationProcessor disposed");
-        while (_transactionChannel.Reader.TryRead(out var transaction))
-            transaction.Fail(disposedException);
-
-        // Cancel any pending invocations waiting for bootstrap responses
-        foreach (var pendingInvocation in _pendingInvocations.Values)
-            pendingInvocation.ResponseTcs.TrySetCanceled(_shutdownCts.Token);
-
-        if (_processingTask != null)
-            try
-            {
-                await _processingTask.ConfigureAwait(false);
-            }
-            catch (OperationCanceledException)
-            {
-                // Expected when task is canceled
-            }
-
-        _state = ProcessorState.Stopped;
-        _shutdownCts.Dispose();
-        _invocationAddedSignal.Dispose();
-    }
-
-    /// <summary>
-    /// Gets a task that completes when initialization finishes (either success or failure).
-    /// Used by LambdaTestServer.StartAsync to wait for initialization to complete.
-    /// </summary>
-    /// <returns>Task that completes with InitResponse indicating success or failure.</returns>
-    internal Task<InitResponse> GetInitCompletionAsync() => _initCompletionTcs.Task;
-
-    /// <summary>
-    /// Creates the HTTP handler for Lambda Bootstrap to use.
-    /// </summary>
-    internal HttpMessageHandler CreateTestingHandler() =>
-        new LambdaTestingHttpHandler(_transactionChannel);
-
-    /// <summary>
-    /// Starts the background processing loop.
-    /// Called by LambdaTestServer after host starts.
-    /// </summary>
-    internal void StartProcessing()
-    {
-        if (_processingTask != null)
-            throw new InvalidOperationException("Processor already started");
-
-        _state = ProcessorState.Initializing;
-        _processingTask = Task.Run(ProcessTransactionsAsync);
-    }
-
-    /// <summary>
-    /// Queues a new invocation to be processed by Lambda Bootstrap.
-    /// Called by LambdaTestServer.InvokeAsync() or LambdaClient.InvokeAsync().
-    /// </summary>
-    internal async Task<InvocationCompletion> QueueInvocationAsync(
-        string requestId,
-        HttpResponseMessage eventResponse,
-        DateTimeOffset deadlineUtc,
-        CancellationToken cancellationToken
-    )
-    {
-        if (_state != ProcessorState.Running)
-            throw new InvalidOperationException(
-                $"Cannot queue invocation when processor is in {_state} state. Processor must be in Running state."
-            );
-
-        var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
-
-        if (!_pendingInvocations.TryAdd(requestId, pending))
-            throw new InvalidOperationException($"Duplicate request ID: {requestId}");
-
-        _pendingInvocationIds.Enqueue(requestId);
-
-        // Signal any waiting /next request that work is available
-        _invocationAddedSignal.Release();
-
-        using var cancellationRegistration = cancellationToken.Register(() =>
-            CancelPendingInvocation(requestId, cancellationToken)
-        );
-
-        // Wait for Bootstrap to process and respond
-        return await pending.ResponseTcs.Task.WaitAsync(cancellationToken);
-    }
-
-    /// <summary>
-    /// Background loop that processes transactions from the handler.
-    /// </summary>
-    private async Task ProcessTransactionsAsync()
-    {
-        await foreach (
-            var transaction in _transactionChannel.Reader.ReadAllAsync(_shutdownCts.Token)
-        )
-            try
-            {
-                await ProcessTransactionAsync(transaction);
-            }
-            catch (Exception ex)
-            {
-                // Fail the transaction and continue processing
-                transaction.Fail(ex);
-            }
-    }
-
-    /// <summary>
-    /// Routes a single transaction based on request type.
-    /// </summary>
-    private async Task ProcessTransactionAsync(LambdaHttpTransaction transaction)
-    {
-        if (!_routeManager.TryMatch(transaction.Request, out var requestType, out var routeValues))
-            throw new InvalidOperationException(
-                $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
-            );
-
-        switch (requestType!.Value)
-        {
-            case RequestType.GetNextInvocation:
-                await HandleGetNextInvocationAsync(transaction);
-                break;
-
-            case RequestType.PostResponse:
-                await HandlePostResponseAsync(transaction, routeValues!);
-                break;
-
-            case RequestType.PostError:
-                await HandlePostErrorAsync(transaction, routeValues!);
-                break;
-
-            case RequestType.PostInitError:
-                await HandlePostInitErrorAsync(transaction);
-                break;
-
-            default:
-                throw new InvalidOperationException(
-                    $"Unexpected request type {requestType} for {transaction.Request.RequestUri}"
-                );
-        }
-    }
-
-    /// <summary>
-    /// Handles GET /invocation/next - Bootstrap polling for work.
-    /// </summary>
-    private async Task HandleGetNextInvocationAsync(LambdaHttpTransaction transaction)
-    {
-        // First successful /next call indicates initialization succeeded
-        if (_state == ProcessorState.Initializing)
-        {
-            _state = ProcessorState.Running;
-            _initCompletionTcs.TrySetResult(
-                new InitResponse { InitStatus = InitStatus.InitCompleted }
-            );
-        }
-
-        // Loop until we find work or shutdown
-        while (true)
-        {
-            // Check if processor is still in a valid state to serve /next requests
-            if (_state != ProcessorState.Running && _state != ProcessorState.Initializing)
-            {
-                transaction.Fail(
-                    new InvalidOperationException(
-                        $"Processor is in {_state} state, cannot serve /next request"
-                    )
-                );
-                return;
-            }
-
-            // Try to dequeue next pending invocation (FIFO)
-            if (TryDequeuePendingInvocation(out var pending))
-            {
-                RespondToNextRequest(transaction, pending);
-                return;
-            }
-
-            // No work available - wait for new invocation or shutdown
-            try
-            {
-                await _invocationAddedSignal.WaitAsync(_shutdownCts.Token);
-                // Loop back to check state and try dequeue again
-            }
-            catch (OperationCanceledException)
-            {
-                // Server shutting down
-                transaction.Fail(
-                    new OperationCanceledException("Server stopped while waiting for work")
-                );
-                return;
-            }
-        }
-    }
-
-    /// <summary>
-    /// Responds to a /next request with a pending invocation.
-    /// </summary>
-    private void RespondToNextRequest(LambdaHttpTransaction transaction, PendingInvocation pending)
-    {
-        // Respond with the event payload and Lambda headers
-        if (transaction.Respond(pending.EventResponse))
-            return;
-
-        // Request was already canceled; re-enqueue invocation to avoid dropping it
-        _pendingInvocationIds.Enqueue(pending.RequestId);
-    }
-
-    /// <summary>
-    /// Handles POST /invocation/{requestId}/response - successful function execution.
-    /// </summary>
-    private async Task HandlePostResponseAsync(
-        LambdaHttpTransaction transaction,
-        RouteValueDictionary routeValues
-    )
-    {
-        var requestId = routeValues["requestId"]?.ToString();
-
-        if (
-            string.IsNullOrEmpty(requestId)
-            || !_pendingInvocations.TryRemove(requestId, out var pending)
-        )
-        {
-            transaction.Respond(
-                new HttpResponseMessage(HttpStatusCode.NotFound)
-                {
-                    Content = new StringContent(
-                        string.IsNullOrEmpty(requestId)
-                            ? "Missing requestId"
-                            : $"No pending invocation for request ID: {requestId}"
-                    ),
-                }
-            );
-            return;
-        }
-
-        // Complete the invocation with the response from Bootstrap
-        pending.ResponseTcs.SetResult(
-            await CreateCompletionAsync(RequestType.PostResponse, transaction.Request)
-        );
-
-        // Acknowledge to Bootstrap
-        transaction.Respond(
-            new HttpResponseMessage(HttpStatusCode.Accepted)
-            {
-                Content = new StringContent(
-                    JsonSerializer.Serialize(
-                        new Dictionary<string, string> { ["status"] = "success" },
-                        _jsonSerializerOptions
-                    ),
-                    Encoding.UTF8,
-                    "application/json"
-                ),
-                Version = Version.Parse("1.1"),
-            }
-        );
-    }
-
-    /// <summary>
-    /// Handles POST /invocation/{requestId}/error - function execution failed.
-    /// </summary>
-    private async Task HandlePostErrorAsync(
-        LambdaHttpTransaction transaction,
-        RouteValueDictionary routeValues
-    )
-    {
-        var requestId = routeValues["requestId"]?.ToString();
-
-        if (
-            string.IsNullOrEmpty(requestId)
-            || !_pendingInvocations.TryRemove(requestId, out var pending)
-        )
-        {
-            transaction.Respond(
-                new HttpResponseMessage(HttpStatusCode.NotFound)
-                {
-                    Content = new StringContent(
-                        string.IsNullOrEmpty(requestId)
-                            ? "Missing requestId"
-                            : $"No pending invocation for request ID: {requestId}"
-                    ),
-                }
-            );
-            return;
-        }
-
-        // Complete the invocation with the error response from Bootstrap
-        pending.ResponseTcs.SetResult(
-            await CreateCompletionAsync(RequestType.PostError, transaction.Request)
-        );
-
-        // Acknowledge to Bootstrap
-        transaction.Respond(
-            new HttpResponseMessage(HttpStatusCode.Accepted)
-            {
-                Content = new StringContent(
-                    "{\"status\":\"success\"}",
-                    Encoding.UTF8,
-                    "application/json"
-                ),
-                Version = Version.Parse("1.1"),
-            }
-        );
-    }
-
-    /// <summary>
-    /// Handles POST /runtime/init/error - Lambda initialization failed.
-    /// Captures the error details and marks initialization as failed.
-    /// The server will be automatically stopped after this is reported.
-    /// </summary>
-    private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction)
-    {
-        // Parse error details from request body
-        ErrorResponse? errorResponse = null;
-        if (transaction.Request.Content != null)
-            try
-            {
-                errorResponse = await transaction.Request.Content.ReadFromJsonAsync<ErrorResponse>(
-                    _jsonSerializerOptions
-                );
-            }
-            catch
-            {
-                // Fallback: create basic error from body text
-                var body = await transaction.Request.Content.ReadAsStringAsync();
-                errorResponse = new ErrorResponse
-                {
-                    ErrorType = transaction.Request.Headers.TryGetValues(
-                        "Lambda-Runtime-Function-Error-Type",
-                        out var values
-                    )
-                        ? values.FirstOrDefault() ?? "Unknown"
-                        : "Unknown",
-                    ErrorMessage = body,
-                };
-            }
-
-        // Mark initialization as failed (only once)
-        if (_state == ProcessorState.Initializing)
-        {
-            _state = ProcessorState.Stopped;
-            _initCompletionTcs.TrySetResult(
-                new InitResponse { InitStatus = InitStatus.InitError, Error = errorResponse }
-            );
-            // Wake up any waiting /next requests by releasing the semaphore
-            // They will check state and fail appropriately
-            _invocationAddedSignal.Release();
-        }
-
-        // Acknowledge to Bootstrap with 202 Accepted
-        transaction.Respond(
-            new HttpResponseMessage(HttpStatusCode.Accepted)
-            {
-                Content = new StringContent(
-                    JsonSerializer.Serialize(
-                        new Dictionary<string, string> { ["status"] = "OK" },
-                        _jsonSerializerOptions
-                    ),
-                    Encoding.UTF8,
-                    "application/json"
-                ),
-                Version = Version.Parse("1.1"),
-            }
-        );
-    }
-
-    private bool TryDequeuePendingInvocation(out PendingInvocation pendingInvocation)
-    {
-        var now = DateTimeOffset.UtcNow;
-
-        while (_pendingInvocationIds.TryDequeue(out var requestId))
-            if (_pendingInvocations.TryGetValue(requestId, out pendingInvocation))
-            {
-                if (pendingInvocation.DeadlineUtc <= now)
-                {
-                    if (_pendingInvocations.TryRemove(requestId, out var expiredInvocation))
-                        expiredInvocation.ResponseTcs.TrySetCanceled();
-
-                    continue;
-                }
-
-                return true;
-            }
-
-        pendingInvocation = null!;
-        return false;
-    }
-
-    private void CancelPendingInvocation(string requestId, CancellationToken cancellationToken)
-    {
-        if (_pendingInvocations.TryRemove(requestId, out var pendingInvocation))
-            pendingInvocation.ResponseTcs.TrySetCanceled(cancellationToken);
-    }
-
-    private static async Task<InvocationCompletion> CreateCompletionAsync(
-        RequestType requestType,
-        HttpRequestMessage sourceRequest
-    )
-    {
-        var clonedRequest = new HttpRequestMessage(sourceRequest.Method, sourceRequest.RequestUri)
-        {
-            Version = sourceRequest.Version,
-            VersionPolicy = sourceRequest.VersionPolicy,
-        };
-
-        foreach (var header in sourceRequest.Headers)
-            clonedRequest.Headers.TryAddWithoutValidation(header.Key, header.Value);
-
-        foreach (var option in sourceRequest.Options)
-            clonedRequest.Options.TryAdd(option.Key, option.Value);
-
-        if (sourceRequest.Content != null)
-        {
-            var contentBytes = await sourceRequest.Content.ReadAsByteArrayAsync();
-            var clonedContent = new ByteArrayContent(contentBytes);
-
-            foreach (var header in sourceRequest.Content.Headers)
-                clonedContent.Headers.TryAddWithoutValidation(header.Key, header.Value);
-
-            clonedRequest.Content = clonedContent;
-        }
-
-        return new InvocationCompletion { Request = clonedRequest, RequestType = requestType };
-    }
-
-    internal void FailPendingInvocations(Exception exception)
-    {
-        while (_transactionChannel.Reader.TryRead(out var transaction))
-            transaction.Fail(exception);
-
-        foreach (var pendingInvocation in _pendingInvocations.Values)
-            pendingInvocation.ResponseTcs.TrySetException(exception);
-    }
-}
diff --git a/src/AwsLambda.Host.Testing/Internal/ProcessorState.cs b/src/AwsLambda.Host.Testing/Internal/ProcessorState.cs
deleted file mode 100644
index 35ad2269..00000000
--- a/src/AwsLambda.Host.Testing/Internal/ProcessorState.cs
+++ /dev/null
@@ -1,35 +0,0 @@
-namespace AwsLambda.Host.Testing;
-
-/// <summary>
-/// Represents the lifecycle state of the InvocationProcessor.
-/// </summary>
-internal enum ProcessorState
-{
-    /// <summary>
-    /// Processor has been created but not yet started.
-    /// </summary>
-    Created,
-
-    /// <summary>
-    /// Processor is running and waiting for initialization to complete.
-    /// Transitions to Running on first successful /next request or to Stopped on init error.
-    /// </summary>
-    Initializing,
-
-    /// <summary>
-    /// Processor is running and processing invocations.
-    /// Initialization completed successfully.
-    /// </summary>
-    Running,
-
-    /// <summary>
-    /// Processor is shutting down.
-    /// Transitions to Stopped when shutdown completes.
-    /// </summary>
-    Stopping,
-
-    /// <summary>
-    /// Processor has stopped and is no longer processing requests.
-    /// </summary>
-    Stopped,
-}
diff --git a/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs b/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs
index 4ce26746..0f2fd484 100644
--- a/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs
+++ b/src/AwsLambda.Host.Testing/LambdaRuntimeRouteManager.cs
@@ -4,7 +4,7 @@
 
 namespace AwsLambda.Host.Testing;
 
-internal class LambdaRuntimeRouteManager : ILambdaRuntimeRouteManager
+internal class LambdaRuntimeRouteManager
 {
     private static readonly RouteTemplate[] Routes =
     [
diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
index 72fce135..f301dc0d 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaServerV2.cs
@@ -50,7 +50,7 @@ public class LambdaServerV2 : IAsyncDisposable
     /// <summary>
     /// Route manager to determine the route of the incoming request from the Lambda.
     /// </summary>
-    private readonly ILambdaRuntimeRouteManager _routeManager = new LambdaRuntimeRouteManager();
+    private readonly LambdaRuntimeRouteManager _routeManager = new();
 
     /// <summary>
     /// CTS used to signal shutdown of the server and cancellation of pending tasks.
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServer.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
deleted file mode 100644
index d5d0fdfd..00000000
--- a/src/AwsLambda.Host.Testing/LambdaTestServer.cs
+++ /dev/null
@@ -1,316 +0,0 @@
-using System.Net;
-using System.Net.Http.Json;
-using System.Text;
-using System.Text.Json;
-using Microsoft.Extensions.Hosting;
-
-namespace AwsLambda.Host.Testing;
-
-/// <summary>
-/// Test server that manages the Lambda host lifecycle and invocation processing.
-/// Provides explicit StartAsync/StopAsync control and direct invocation capabilities.
-/// </summary>
-public class LambdaTestServer : IAsyncDisposable
-{
-    private readonly IHost _host;
-    private readonly InvocationProcessor _processor;
-    private readonly LambdaClientOptions _defaultClientOptions;
-    private readonly JsonSerializerOptions _jsonSerializerOptions;
-    private readonly Task<Exception?> _entryPointCompletion;
-    private int _requestCounter;
-    private ServerState _state;
-
-    internal LambdaTestServer(
-        IHost host,
-        InvocationProcessor processor,
-        LambdaClientOptions? defaultClientOptions = null,
-        JsonSerializerOptions? jsonSerializerOptions = null,
-        Task<Exception?>? entryPointCompletion = null
-    )
-    {
-        _host = host ?? throw new ArgumentNullException(nameof(host));
-        _processor = processor ?? throw new ArgumentNullException(nameof(processor));
-        _defaultClientOptions = defaultClientOptions ?? new LambdaClientOptions();
-        _jsonSerializerOptions = jsonSerializerOptions ?? new JsonSerializerOptions();
-        _entryPointCompletion = entryPointCompletion ?? Task.FromResult<Exception?>(null);
-        _state = ServerState.Created;
-    }
-
-    /// <summary>
-    /// Gets the current lifecycle state of the server.
-    /// </summary>
-    public ServerState State => _state;
-
-    /// <summary>
-    /// Gets the service provider from the host.
-    /// Only available after the server has been started.
-    /// </summary>
-    /// <exception cref="InvalidOperationException">Thrown if server is not in Running state.</exception>
-    public IServiceProvider Services => _host.Services;
-
-    /// <summary>
-    /// Starts the Lambda host and waits for initialization to complete.
-    /// Returns initialization result indicating success or failure.
-    /// If initialization fails, the server is automatically stopped.
-    /// </summary>
-    /// <param name="cancellationToken">Cancellation token.</param>
-    /// <returns>InitResponse indicating whether initialization succeeded or failed.</returns>
-    /// <exception cref="InvalidOperationException">Thrown if server is not in Created state.</exception>
-    public async Task<InitResponse> StartAsync(CancellationToken cancellationToken = default)
-    {
-        if (_state != ServerState.Created)
-            throw new InvalidOperationException(
-                $"Server can only be started from Created state. Current state: {_state}"
-            );
-
-        try
-        {
-            _state = ServerState.Starting;
-
-            // Start the host
-            await _host.StartAsync(cancellationToken);
-
-            // Start background processing
-            _processor.StartProcessing();
-
-            _state = ServerState.Running;
-
-            // Race initialization completion against entry point failure
-            var initTask = _processor.GetInitCompletionAsync();
-            var completed = await Task.WhenAny(initTask, _entryPointCompletion)
-                .WaitAsync(cancellationToken);
-
-            // Check if entry point failed before init completed
-            if (completed == _entryPointCompletion)
-            {
-                var exception = await _entryPointCompletion;
-                if (exception != null)
-                    throw new InvalidOperationException(
-                        "Entry point failed during initialization",
-                        exception
-                    );
-            }
-
-            // Wait for init to complete
-            var initResponse = await initTask;
-
-            // If init failed, auto-stop the server
-            // if (initResponse != InitStatus.InitCompleted)
-            //     await StopAsync(CancellationToken.None);
-            // if (!initResponse.InitSuccess)
-            //     await StopAsync(CancellationToken.None);
-
-            return initResponse;
-        }
-        catch
-        {
-            // Cleanup on failure
-            await DisposeAsync();
-            throw;
-        }
-    }
-
-    /// <summary>
-    /// Stops the Lambda host.
-    /// </summary>
-    /// <param name="cancellationToken">Cancellation token.</param>
-    public async Task StopAsync(CancellationToken cancellationToken = default)
-    {
-        // Idempotent - allow multiple calls
-        if (_state == ServerState.Stopped || _state == ServerState.Disposed)
-            return;
-
-        if (_state != ServerState.Running)
-            throw new InvalidOperationException(
-                $"Server can only be stopped from Running state. Current state: {_state}"
-            );
-
-        try
-        {
-            _state = ServerState.Stopping;
-            await _host.StopAsync(cancellationToken);
-            _state = ServerState.Stopped;
-        }
-        catch (OperationCanceledException)
-        {
-            // Still mark as stopped even if shutdown was canceled
-            _state = ServerState.Stopped;
-            throw;
-        }
-    }
-
-    /// <summary>
-    /// Invokes the Lambda function with the given event and waits for the response.
-    /// </summary>
-    /// <typeparam name="TResponse">The type of the response.</typeparam>
-    /// <typeparam name="TEvent">The type of the event.</typeparam>
-    /// <param name="invokeEvent">The event to invoke with.</param>
-    /// <param name="cancellationToken">Cancellation token.</param>
-    /// <returns>The invocation response.</returns>
-    /// <exception cref="InvalidOperationException">Thrown if server is not in Running state.</exception>
-    public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
-        TEvent invokeEvent,
-        CancellationToken cancellationToken = default
-    ) => await InvokeAsync<TResponse, TEvent>(invokeEvent, null, cancellationToken);
-
-    /// <summary>
-    /// Invokes the Lambda function with the given event and waits for the response.
-    /// </summary>
-    /// <typeparam name="TResponse">The type of the response.</typeparam>
-    /// <typeparam name="TEvent">The type of the event.</typeparam>
-    /// <param name="invokeEvent">The event to invoke with.</param>
-    /// <param name="options">Client options for this invocation.</param>
-    /// <param name="cancellationToken">Cancellation token.</param>
-    /// <returns>The invocation response.</returns>
-    /// <exception cref="InvalidOperationException">Thrown if server is not in Running state.</exception>
-    public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
-        TEvent invokeEvent,
-        LambdaClientOptions? options,
-        CancellationToken cancellationToken = default
-    )
-    {
-        if (_state != ServerState.Running)
-            throw new InvalidOperationException(
-                _state == ServerState.Created
-                    ? "Server must be started before invoking. Call StartAsync() first."
-                    : $"Server is {_state}. Only Running servers can process invocations."
-            );
-
-        // Use provided options or fall back to defaults
-        var clientOptions = options ?? _defaultClientOptions;
-
-        // Generate unique request ID
-        var requestId = GetRequestId();
-
-        // Create the event response with Lambda headers
-        var eventResponse = CreateEventResponse(invokeEvent, requestId, clientOptions);
-        var deadlineUtc = DateTimeOffset.UtcNow.Add(
-            clientOptions.InvocationHeaderOptions.ClientWaitTimeout
-        );
-
-        // Queue invocation and wait for Bootstrap to process it
-        var waitTimeout = clientOptions.InvocationHeaderOptions.ClientWaitTimeout;
-
-        using var timeoutCts = new CancellationTokenSource(waitTimeout);
-        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(
-            cancellationToken,
-            timeoutCts.Token
-        );
-
-        var completion = await _processor.QueueInvocationAsync(
-            requestId,
-            eventResponse,
-            deadlineUtc,
-            linkedCts.Token
-        );
-
-        var responseMessage = completion.Request;
-        var wasSuccess = completion.RequestType == RequestType.PostResponse;
-
-        var invocationResponse = new InvocationResponse<TResponse>
-        {
-            WasSuccess = wasSuccess,
-            Response = wasSuccess
-                ? await (
-                    responseMessage.Content?.ReadFromJsonAsync<TResponse>(
-                        _jsonSerializerOptions,
-                        cancellationToken
-                    ) ?? Task.FromResult<TResponse?>(default)
-                )
-                : default,
-            Error = !wasSuccess
-                ? await (
-                    responseMessage.Content?.ReadFromJsonAsync<ErrorResponse>(
-                        _jsonSerializerOptions,
-                        cancellationToken
-                    ) ?? Task.FromResult<ErrorResponse?>(null)
-                )
-                : null,
-        };
-
-        return invocationResponse;
-    }
-
-    /// <summary>
-    /// Disposes the server, stopping the host if needed.
-    /// </summary>
-    public async ValueTask DisposeAsync()
-    {
-        if (_state == ServerState.Disposed)
-            return;
-
-        try
-        {
-            // Stop if running
-            if (_state == ServerState.Running)
-                try
-                {
-                    await StopAsync();
-                }
-                catch
-                {
-                    // Best effort - continue with disposal
-                }
-
-            // Dispose processor
-            await _processor.DisposeAsync();
-
-            // Dispose host
-            if (_host is IAsyncDisposable asyncDisposable)
-                await asyncDisposable.DisposeAsync();
-            else
-                _host.Dispose();
-        }
-        finally
-        {
-            _state = ServerState.Disposed;
-        }
-    }
-
-    private HttpResponseMessage CreateEventResponse<TEvent>(
-        TEvent invokeEvent,
-        string requestId,
-        LambdaClientOptions clientOptions
-    )
-    {
-        var response = new HttpResponseMessage(HttpStatusCode.OK)
-        {
-            Content = new StringContent(
-                JsonSerializer.Serialize(invokeEvent, _jsonSerializerOptions),
-                Encoding.UTF8,
-                "application/json"
-            ),
-            Version = Version.Parse("1.1"),
-        };
-
-        // Add standard HTTP headers
-        response.Headers.Date = new DateTimeOffset(clientOptions.InvocationHeaderOptions.Date);
-        response.Headers.TransferEncodingChunked = clientOptions
-            .InvocationHeaderOptions
-            .TransferEncodingChunked;
-
-        // Add custom Lambda runtime headers
-        var deadlineMs = DateTimeOffset
-            .UtcNow.Add(clientOptions.InvocationHeaderOptions.FunctionTimeout)
-            .ToUnixTimeMilliseconds();
-        response.Headers.Add("Lambda-Runtime-Deadline-Ms", deadlineMs.ToString());
-        response.Headers.Add("Lambda-Runtime-Aws-Request-Id", requestId);
-        response.Headers.Add(
-            "Lambda-Runtime-Trace-Id",
-            clientOptions.InvocationHeaderOptions.TraceId
-        );
-        response.Headers.Add(
-            "Lambda-Runtime-Invoked-Function-Arn",
-            clientOptions.InvocationHeaderOptions.FunctionArn
-        );
-
-        // Add any additional custom headers
-        foreach (var header in clientOptions.InvocationHeaderOptions.AdditionalHeaders)
-            response.Headers.Add(header.Key, header.Value);
-
-        return response;
-    }
-
-    private string GetRequestId() =>
-        Interlocked.Increment(ref _requestCounter).ToString().PadLeft(12, '0');
-}
diff --git a/src/AwsLambda.Host.Testing/Models/InitResponse.cs b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
index c15b754f..833a22b2 100644
--- a/src/AwsLambda.Host.Testing/Models/InitResponse.cs
+++ b/src/AwsLambda.Host.Testing/Models/InitResponse.cs
@@ -15,24 +15,3 @@ public class InitResponse
     /// </summary>
     public InitStatus InitStatus { get; internal init; }
 }
-
-/// <summary>
-/// An enumeration of possible statuses for Lambda initialization.
-/// </summary>
-public enum InitStatus
-{
-    /// <summary>
-    /// Initialization of the Lambda completed successfully.
-    /// </summary>
-    InitCompleted,
-
-    /// <summary>
-    /// Initialization of the Lambda failed, and the Lambda returned an error.
-    /// </summary>
-    InitError,
-
-    /// <summary>
-    /// Initialization of the Lambda failed, and the Host process exited.
-    /// </summary>
-    HostExited,
-}
diff --git a/src/AwsLambda.Host.Testing/Models/InitStatus.cs b/src/AwsLambda.Host.Testing/Models/InitStatus.cs
new file mode 100644
index 00000000..eee1c0ff
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/Models/InitStatus.cs
@@ -0,0 +1,22 @@
+namespace AwsLambda.Host.Testing;
+
+/// <summary>
+/// An enumeration of possible statuses for Lambda initialization.
+/// </summary>
+public enum InitStatus
+{
+    /// <summary>
+    /// Initialization of the Lambda completed successfully.
+    /// </summary>
+    InitCompleted,
+
+    /// <summary>
+    /// Initialization of the Lambda failed, and the Lambda returned an error.
+    /// </summary>
+    InitError,
+
+    /// <summary>
+    /// Initialization of the Lambda failed, and the Host process exited.
+    /// </summary>
+    HostExited,
+}
diff --git a/src/AwsLambda.Host.Testing/Models/LambdaBootstrapRequest.cs b/src/AwsLambda.Host.Testing/Models/LambdaBootstrapRequest.cs
deleted file mode 100644
index b4ab21da..00000000
--- a/src/AwsLambda.Host.Testing/Models/LambdaBootstrapRequest.cs
+++ /dev/null
@@ -1,24 +0,0 @@
-using Microsoft.AspNetCore.Routing;
-
-namespace AwsLambda.Host.Testing;
-
-internal class LambdaBootstrapRequest
-{
-    internal string? RequestId
-    {
-        get
-        {
-            field ??=
-                RouteValue.TryGetValue("RequestId", out var requestId)
-                && requestId is string requestIdString
-                    ? requestIdString
-                    : null;
-
-            return field;
-        }
-    }
-
-    internal required HttpRequestMessage RequestMessage { get; init; }
-    internal required RequestType RequestType { get; init; }
-    internal required RouteValueDictionary RouteValue { get; init; }
-}
diff --git a/src/AwsLambda.Host.Testing/ServerState.cs b/src/AwsLambda.Host.Testing/Models/ServerState.cs
similarity index 100%
rename from src/AwsLambda.Host.Testing/ServerState.cs
rename to src/AwsLambda.Host.Testing/Models/ServerState.cs
diff --git a/src/AwsLambda.Host.Testing/runtime-api.yaml b/src/AwsLambda.Host.Testing/runtime-api.yaml
deleted file mode 100644
index 9acfba2d..00000000
--- a/src/AwsLambda.Host.Testing/runtime-api.yaml
+++ /dev/null
@@ -1,250 +0,0 @@
-#
-# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# 
-# Licensed under the Apache License, Version 2.0 (the "License").
-# You may not use this file except in compliance with the License.
-# A copy of the License is located at
-#
-#  http://aws.amazon.com/apache2.0
-#
-# or in the "license" file accompanying this file. This file is distributed
-# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-# express or implied. See the License for the specific language governing
-# permissions and limitations under the License.
-#
-# This file was copied from https://docs.aws.amazon.com/lambda/latest/dg/samples/runtime-api.zip on 12/1/2018
-#
-# This file documents BYOL (RAPID) API using Open API 3 specification format.
-#
-#
-# A note on error reporting in BYOL.
-#
-# Runtimes should report all errors using Lambda standard error format, in order to integrate with other AWS services:
-#
-# Content-Type: application/vnd.aws.lambda.error+json:
-# {
-#     "errorMessage": "...",
-#     "errorType": "...",
-#     "stackTrace": [],
-# }
-#
-# Corresponding Open API 3 schema:
-#
-# ErrorRequest:
-#   type: object
-#   properties:
-#     errorMessage:
-#       type: string
-#     errorType:
-#       type: string
-#     stackTrace:
-#       type: array
-#       items:
-#         type: string
-#
-# Lambda's default behavior is to use Lambda-Runtime-Function-Error-Type header value to construct an error response
-# when error payload is not provided or can not be read.
-
-
-openapi: 3.0.0
-info:
-  title: BYOL API
-  description: Native Runtime API.
-  version: 1.0.2
-
-servers:
-  - url: /2018-06-01
-
-paths:
-
-  /runtime/init/error:
-    post:
-      summary: >
-        Non-recoverable initialization error. Runtime should exit after reporting
-        the error. Error will be served in response to the first invoke.
-      parameters:
-        - in: header
-          name: Lambda-Runtime-Function-Error-Type
-          schema:
-            type: string
-      requestBody:
-        content:
-          '*/*':
-            schema: { }
-      responses:
-        '202':
-          description: Accepted
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/StatusResponse'
-        '403':
-          description: Forbidden
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/ErrorResponse'
-        '500':
-          description: >
-            Container error. Non-recoverable state. Runtime should exit promptly.
-
-  /runtime/invocation/next:
-    get:
-      summary: >
-        Runtime makes this HTTP request when it is ready to receive and process a
-        new invoke.
-      responses:
-        '200':
-          description: >
-            This is an iterator-style blocking API call. Response contains
-            event JSON document, specific to the invoking service.
-          headers:
-            Lambda-Runtime-Aws-Request-Id:
-              description: AWS request ID associated with the request.
-              schema:
-                type: string
-            Lambda-Runtime-Trace-Id:
-              description: X-Ray tracing header.
-              schema:
-                type: string
-            Lambda-Runtime-Client-Context:
-              description: >
-                Information about the client application and device when invoked
-                through the AWS Mobile SDK.
-              schema:
-                type: string
-            Lambda-Runtime-Cognito-Identity:
-              description: >
-                Information about the Amazon Cognito identity provider when invoked
-                through the AWS Mobile SDK.
-              schema:
-                type: string
-            Lambda-Runtime-Deadline-Ms:
-              description: >
-                Function execution deadline counted in milliseconds since the Unix epoch.
-              schema:
-                type: string
-            Lambda-Runtime-Invoked-Function-Arn:
-              description: >
-                The ARN requested. This can be different in each invoke that
-                executes the same version.
-              schema:
-                type: string
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/EventResponse'
-        '403':
-          description: Forbidden
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/ErrorResponse'
-        '500':
-          description: >
-            Container error. Non-recoverable state. Runtime should exit promptly.
-
-  /runtime/invocation/{AwsRequestId}/response:
-    post:
-      summary: Runtime makes this request in order to submit a response.
-      parameters:
-        - in: path
-          name: AwsRequestId
-          schema:
-            type: string
-          required: true
-      requestBody:
-        content:
-          '*/*':
-            schema: { }
-      responses:
-        '202':
-          description: Accepted
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/StatusResponse'
-        '400':
-          description: Bad Request
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/ErrorResponse'
-        '403':
-          description: Forbidden
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/ErrorResponse'
-        '413':
-          description: Payload Too Large
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/ErrorResponse'
-        '500':
-          description: >
-            Container error. Non-recoverable state. Runtime should exit promptly.
-
-  /runtime/invocation/{AwsRequestId}/error:
-    post:
-      summary: >
-        Runtime makes this request in order to submit an error response. It can
-        be either a function error, or a runtime error. Error will be served in
-        response to the invoke.
-      parameters:
-        - in: path
-          name: AwsRequestId
-          schema:
-            type: string
-          required: true
-        - in: header
-          name: Lambda-Runtime-Function-Error-Type
-          schema:
-            type: string
-      requestBody:
-        content:
-          '*/*':
-            schema: { }
-      responses:
-        '202':
-          description: Accepted
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/StatusResponse'
-        '400':
-          description: Bad Request
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/ErrorResponse'
-        '403':
-          description: Forbidden
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/ErrorResponse'
-        '500':
-          description: >
-            Container error. Non-recoverable state. Runtime should exit promptly.
-
-components:
-  schemas:
-
-    StatusResponse:
-      type: object
-      properties:
-        status:
-          type: string
-
-    ErrorResponse:
-      type: object
-      properties:
-        errorMessage:
-          type: string
-        errorType:
-          type: string
-
-    EventResponse:
-      type: object
\ No newline at end of file

From 2d1db369e18c908e4c7e0cc6a6a5be522616b096 Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 11:31:50 -0500
Subject: [PATCH 27/29] refactor(testing): rename LambdaServerV2 to
 LambdaTestServer for clarity

- Replaced all occurrences of `LambdaServerV2` with `LambdaTestServer` across testing infrastructure.
- Updated references in `LambdaApplicationFactory` and related methods to use `LambdaTestServer`.
- Adjusted test cases to align with the updated `LambdaTestServer` naming convention.
- Updated XML documentation and comments to reflect the new naming.
---
 .../Tests/LambdaHostTest.cs                   | 36 +++++++++----------
 .../LambdaApplicationFactory.cs               | 22 ++++++------
 .../LambdaHttpHandler.cs                      |  2 +-
 ...{LambdaServerV2.cs => LambdaTestServer.cs} | 16 ++++-----
 .../Models/ServerState.cs                     | 12 +++----
 5 files changed, 44 insertions(+), 44 deletions(-)
 rename src/AwsLambda.Host.Testing/{LambdaServerV2.cs => LambdaTestServer.cs} (96%)

diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 5d81f73a..9ac64012 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -15,11 +15,11 @@ public async Task LambdaHost_CanStartWithoutError()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
 
-        var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+        var setup = await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
 
         setup.InitStatus.Should().Be(InitStatus.InitCompleted);
 
-        var response = await factory.Server.InvokeAsync<string, string>(
+        var response = await factory.TestServer.InvokeAsync<string, string>(
             "Jonas",
             TestContext.Current.CancellationToken
         );
@@ -34,11 +34,11 @@ public async Task LambdaHost_HandlerReturnsError()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
 
-        var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+        var setup = await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
 
         setup.InitStatus.Should().Be(InitStatus.InitCompleted);
 
-        var response = await factory.Server.InvokeAsync<string, string>(
+        var response = await factory.TestServer.InvokeAsync<string, string>(
             "",
             TestContext.Current.CancellationToken
         );
@@ -53,11 +53,11 @@ public async Task LambdaHost_CanBeShutdown()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
 
-        var setup = await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+        var setup = await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
 
         setup.InitStatus.Should().Be(InitStatus.InitCompleted);
 
-        var response = await factory.Server.InvokeAsync<string, string>(
+        var response = await factory.TestServer.InvokeAsync<string, string>(
             "Jonas",
             TestContext.Current.CancellationToken
         );
@@ -66,7 +66,7 @@ public async Task LambdaHost_CanBeShutdown()
         response.Should().NotBeNull();
         response.Response.Should().Be("Hello Jonas!");
 
-        await factory.Server.StopAsync(TestContext.Current.CancellationToken);
+        await factory.TestServer.StopAsync(TestContext.Current.CancellationToken);
     }
 
     [Fact]
@@ -86,7 +86,7 @@ public async Task LambdaHost_ServerInternalExceptions_AreCaughtAndReturnedAsErro
         });
 
         var act = async () =>
-            await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+            await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
 
         (await act.Should().ThrowAsync<AggregateException>())
             .And.InnerExceptions.Should()
@@ -102,13 +102,13 @@ ex is InvalidOperationException
     public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+        await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
 
         // Launch 5 concurrent invocations
         var tasks = Enumerable
             .Range(1, 5)
             .Select(i =>
-                factory.Server.InvokeAsync<string, string>(
+                factory.TestServer.InvokeAsync<string, string>(
                     $"User{i}",
                     TestContext.Current.CancellationToken
                 )
@@ -130,9 +130,9 @@ public async Task LambdaHost_ProcessesConcurrentInvocationsInFifoOrder()
     public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+        await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
 
-        var response = await factory.Server.InvokeAsync<string, int>(
+        var response = await factory.TestServer.InvokeAsync<string, int>(
             123,
             TestContext.Current.CancellationToken
         );
@@ -146,13 +146,13 @@ public async Task InvokeAsync_WithInvalidPayload_ReturnsError()
     public async Task InvokeAsync_WithPreCanceledToken_CancelsInvocation()
     {
         await using var factory = new LambdaApplicationFactory<Program>();
-        await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+        await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
 
         using var cts = new CancellationTokenSource();
         await cts.CancelAsync();
 
         await Assert.ThrowsAsync<TaskCanceledException>(() =>
-            factory.Server.InvokeAsync<string, string>("Jonas", cts.Token)
+            factory.TestServer.InvokeAsync<string, string>("Jonas", cts.Token)
         );
     }
 
@@ -161,12 +161,12 @@ await Assert.ThrowsAsync<TaskCanceledException>(() =>
     //     await Assert.ThrowsAsync<TaskCanceledException>(async () =>
     //     {
     //         await using var factory = new LambdaApplicationFactory<Program>();
-    //         await factory.Server.StartAsync(TestContext.Current.CancellationToken);
+    //         await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
     //
     //         var options = new LambdaClientOptions();
     //         options.InvocationHeaderOptions.ClientWaitTimeout = TimeSpan.Zero;
     //
-    //         await factory.Server.InvokeAsync<string, string>(
+    //         await factory.TestServer.InvokeAsync<string, string>(
     //             "Jonas",
     //             options,
     //             TestContext.Current.CancellationToken
@@ -182,10 +182,10 @@ await Assert.ThrowsAsync<TaskCanceledException>(() =>
     //     await using var factory = new LambdaApplicationFactory<Program>();
     //
     //     using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
-    //     var initResponse = await factory.Server.StartAsync(cts.Token);
+    //     var initResponse = await factory.TestServer.StartAsync(cts.Token);
     //
     //     Assert.False(initResponse.InitSuccess);
     //     Assert.NotNull(initResponse.Error);
-    //     Assert.Equal(ServerState.Stopped, factory.Server.State);
+    //     Assert.Equal(ServerState.Stopped, factory.TestServer.State);
     // }
 }
diff --git a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
index 2fa533ab..a04f3925 100644
--- a/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
+++ b/src/AwsLambda.Host.Testing/LambdaApplicationFactory.cs
@@ -30,13 +30,13 @@ public class LambdaApplicationFactory<TEntryPoint> : IDisposable, IAsyncDisposab
     private bool _disposed;
     private bool _disposedAsync;
     private IHost? _host;
-    private LambdaServerV2? _server;
+    private LambdaTestServer? _server;
 
     /// <summary>
     /// <para>
     /// Creates an instance of <see cref="LambdaApplicationFactory{TEntryPoint}"/>. This factory can be used to
-    /// create a <see cref="LambdaServerV2"/> instance using the MVC application defined by <typeparamref name="TEntryPoint"/>
-    /// and one or more <see cref="HttpClient"/> instances used to send <see cref="HttpRequestMessage"/> to the <see cref="LambdaServerV2"/>.
+    /// create a <see cref="LambdaTestServer"/> instance using the MVC application defined by <typeparamref name="TEntryPoint"/>
+    /// and one or more <see cref="HttpClient"/> instances used to send <see cref="HttpRequestMessage"/> to the <see cref="LambdaTestServer"/>.
     /// The <see cref="LambdaApplicationFactory{TEntryPoint}"/> will find the entry point class of <typeparamref name="TEntryPoint"/>
     /// assembly and initialize the application by calling <c>IHostBuilder CreateWebHostBuilder(string [] args)</c>
     /// on <typeparamref name="TEntryPoint"/>.
@@ -71,9 +71,9 @@ public class LambdaApplicationFactory<TEntryPoint> : IDisposable, IAsyncDisposab
         _derivedFactories.AsReadOnly();
 
     /// <summary>
-    /// Gets the <see cref="LambdaServerV2"/> created by this <see cref="LambdaApplicationFactory{TEntryPoint}"/>.
+    /// Gets the <see cref="LambdaTestServer"/> created by this <see cref="LambdaApplicationFactory{TEntryPoint}"/>.
     /// </summary>
-    public LambdaServerV2 Server
+    public LambdaTestServer TestServer
     {
         get
         {
@@ -85,7 +85,7 @@ public LambdaServerV2 Server
     /// <summary>
     /// Gets the <see cref="IServiceProvider"/> created by the server associated with this <see cref="LambdaApplicationFactory{TEntryPoint}"/>.
     /// </summary>
-    public virtual IServiceProvider Services => Server.Services;
+    public virtual IServiceProvider Services => TestServer.Services;
 
     /// <inheritdoc />
     public virtual async ValueTask DisposeAsync()
@@ -99,7 +99,7 @@ public virtual async ValueTask DisposeAsync()
         foreach (var factory in _derivedFactories)
             await ((IAsyncDisposable)factory).DisposeAsync().ConfigureAwait(false);
 
-        // Server handles disposing both processor and host
+        // TestServer handles disposing both processor and host
         if (_server != null)
             await _server.DisposeAsync().ConfigureAwait(false);
 
@@ -213,7 +213,7 @@ private void ConfigureHostBuilder(
         SetContentRoot(hostBuilder);
         _configuration(hostBuilder);
 
-        _server = new LambdaServerV2(entryPointCompletion);
+        _server = new LambdaTestServer(entryPointCompletion);
 
         // set Lambda Bootstrap Http Client
         hostBuilder.ConfigureServices(services =>
@@ -232,7 +232,7 @@ private void ConfigureHostBuilder(
         _server.SetHost(_host);
 
         // Create the public server with the built (but not started) host
-        // _server = new LambdaServerV2(_host, processor, entryPointCompletion:
+        // _server = new LambdaTestServer(_host, processor, entryPointCompletion:
         // entryPointCompletion);
     }
 
@@ -400,7 +400,7 @@ private static void EnsureDepsFile()
     }
 
     /// <summary>
-    /// Creates a <see cref="IHostBuilder"/> used to set up <see cref="LambdaServerV2"/>.
+    /// Creates a <see cref="IHostBuilder"/> used to set up <see cref="LambdaTestServer"/>.
     /// </summary>
     /// <remarks>
     /// The default implementation of this method looks for a <c>public static IHostBuilder CreateHostBuilder(string[] args)</c>
@@ -426,7 +426,7 @@ private static void EnsureDepsFile()
     /// <param name="builder">The <see cref="IHostBuilder"/> used to create the host.</param>
     /// <returns>The <see cref="IHost"/> with the bootstrapped application.</returns>
     protected virtual IHost CreateHost(IHostBuilder builder) =>
-        // Build the host but DON'T start it - LambdaServerV2.StartAsync() will start it
+        // Build the host but DON'T start it - LambdaTestServer.StartAsync() will start it
         builder.Build();
 
     /// <summary>
diff --git a/src/AwsLambda.Host.Testing/LambdaHttpHandler.cs b/src/AwsLambda.Host.Testing/LambdaHttpHandler.cs
index bf7cc1a9..ffc9f4df 100644
--- a/src/AwsLambda.Host.Testing/LambdaHttpHandler.cs
+++ b/src/AwsLambda.Host.Testing/LambdaHttpHandler.cs
@@ -41,7 +41,7 @@ CancellationToken cancellationToken
         }
         catch (ChannelClosedException)
         {
-            // Server is shutting down; propagate cancellation to caller
+            // TestServer is shutting down; propagate cancellation to caller
             var canceled = new TaskCompletionSource<HttpResponseMessage>(
                 TaskCreationOptions.RunContinuationsAsynchronously
             );
diff --git a/src/AwsLambda.Host.Testing/LambdaServerV2.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
similarity index 96%
rename from src/AwsLambda.Host.Testing/LambdaServerV2.cs
rename to src/AwsLambda.Host.Testing/LambdaTestServer.cs
index f301dc0d..46d96ad3 100644
--- a/src/AwsLambda.Host.Testing/LambdaServerV2.cs
+++ b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
@@ -11,7 +11,7 @@
 
 namespace AwsLambda.Host.Testing;
 
-public class LambdaServerV2 : IAsyncDisposable
+public class LambdaTestServer : IAsyncDisposable
 {
     /// <summary>
     /// Options used to configure how the server interacts with the Lambda.
@@ -90,7 +90,7 @@ public class LambdaServerV2 : IAsyncDisposable
     /// </summary>
     private ServerState _state;
 
-    internal LambdaServerV2(
+    internal LambdaTestServer(
         Task<Exception?>? entryPointCompletion,
         CancellationToken shutdownToken = default
     )
@@ -136,7 +136,7 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
     {
         if (_state != ServerState.Created)
             throw new InvalidOperationException(
-                "Server has already been started and cannot be restarted."
+                "TestServer has already been started and cannot be restarted."
             );
 
         if (_host is null)
@@ -175,7 +175,7 @@ await TaskHelpers
         }
 
         throw new InvalidOperationException(
-            "Server initialization failed with neither an error nor completion."
+            "TestServer initialization failed with neither an error nor completion."
         );
     }
 
@@ -186,7 +186,7 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
     {
         if (_state != ServerState.Running)
             throw new InvalidOperationException(
-                "Server is not Running and as such an event cannot be invoked."
+                "TestServer is not Running and as such an event cannot be invoked."
             );
 
         using var cts = CancellationTokenSource.CreateLinkedTokenSource(
@@ -246,7 +246,7 @@ public async Task StopAsync(CancellationToken cancellationToken = default)
     {
         if (_state != ServerState.Running)
             throw new InvalidOperationException(
-                "Server is not running and as such cannot be stopped."
+                "TestServer is not running and as such cannot be stopped."
             );
 
         _state = ServerState.Stopping;
@@ -264,7 +264,7 @@ await TaskHelpers
     }
 
     //      ┌──────────────────────────────────────────────────────────┐
-    //      │                  Internal Server Logic                   │
+    //      │                  Internal TestServer Logic                   │
     //      └──────────────────────────────────────────────────────────┘
 
     private async Task ProcessTransactionsAsync()
@@ -378,7 +378,7 @@ private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction)
             );
 
         throw new InvalidOperationException(
-            "Server is already started and as such an initialization error cannot be reported."
+            "TestServer is already started and as such an initialization error cannot be reported."
         );
     }
 
diff --git a/src/AwsLambda.Host.Testing/Models/ServerState.cs b/src/AwsLambda.Host.Testing/Models/ServerState.cs
index 9e2f394e..892b610d 100644
--- a/src/AwsLambda.Host.Testing/Models/ServerState.cs
+++ b/src/AwsLambda.Host.Testing/Models/ServerState.cs
@@ -6,32 +6,32 @@ namespace AwsLambda.Host.Testing;
 public enum ServerState
 {
     /// <summary>
-    /// Server created but not started.
+    /// TestServer created but not started.
     /// </summary>
     Created,
 
     /// <summary>
-    /// Server is starting (building host).
+    /// TestServer is starting (building host).
     /// </summary>
     Starting,
 
     /// <summary>
-    /// Server is running and accepting invocations.
+    /// TestServer is running and accepting invocations.
     /// </summary>
     Running,
 
     /// <summary>
-    /// Server is stopping.
+    /// TestServer is stopping.
     /// </summary>
     Stopping,
 
     /// <summary>
-    /// Server has stopped cleanly.
+    /// TestServer has stopped cleanly.
     /// </summary>
     Stopped,
 
     /// <summary>
-    /// Server has been disposed.
+    /// TestServer has been disposed.
     /// </summary>
     Disposed,
 }

From 7cca257827e0f45fb55318cf461c8942937dcf2d Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 15:27:28 -0500
Subject: [PATCH 28/29] feat(testing): update Lambda server options and
 invocation handling

- Renamed `LambdaClientOptions` to `LambdaServerOptions` and refactored related usage.
- Updated invocation timeout handling to use `FunctionTimeout` in `LambdaServerOptions`.
- Refined `InvokeAsync` method to include `traceId` and improved error handling for pending IDs.
- Enhanced response creation with additional headers and updated `traceId` logic.
- Removed unused properties like `Date`, `InvocationTimeout`, and `TransferEncodingChunked`.
- Improved cancellation logic with `LinkedCtsWithInvocationDeadline` and `LinkedCts` methods.
- Introduced `LambdaTestServerExtensions` for simplified invocation in test cases.
---
 .../Lambda/Program.cs                         |  5 +-
 .../Tests/LambdaHostTest.cs                   |  7 +-
 .../LambdaTestServer.cs                       | 74 ++++++++++---------
 .../LambdaTestServerExtensions.cs             | 16 ++++
 .../Options/LambdaClientOptions.cs            | 12 ---
 ...rOptions.cs => LambdaInvocationOptions.cs} | 31 +-------
 .../Options/LambdaServerOptions.cs            | 28 +++++++
 7 files changed, 92 insertions(+), 81 deletions(-)
 create mode 100644 src/AwsLambda.Host.Testing/LambdaTestServerExtensions.cs
 delete mode 100644 src/AwsLambda.Host.Testing/Options/LambdaClientOptions.cs
 rename src/AwsLambda.Host.Testing/Options/{LambdaInvocationHeaderOptions.cs => LambdaInvocationOptions.cs} (54%)
 create mode 100644 src/AwsLambda.Host.Testing/Options/LambdaServerOptions.cs

diff --git a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
index 59f49a23..e641ddc1 100644
--- a/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Lambda/Program.cs
@@ -1,4 +1,5 @@
 ﻿using AwsLambda.Host.Builder;
+using AwsLambda.Host.Core;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 
@@ -23,9 +24,9 @@
 
 // Map your handler - the event is automatically injected
 lambda.MapHandler(
-    async ([Event] string name, CancellationToken cancellationToken) =>
+    async ([Event] string name, ILambdaHostContext context, CancellationToken cancellationToken) =>
     {
-        // await Task.Delay(TimeSpan.FromSeconds(60), cancellationToken);
+        await Task.Delay(TimeSpan.FromSeconds(60), cancellationToken);
         if (string.IsNullOrWhiteSpace(name))
             throw new ArgumentNullException(nameof(name), "Name is required.");
 
diff --git a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
index 9ac64012..81a1e338 100644
--- a/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
+++ b/examples/AwsLambda.Host.Example.Testing/Tests/LambdaHostTest.cs
@@ -21,6 +21,7 @@ public async Task LambdaHost_CanStartWithoutError()
 
         var response = await factory.TestServer.InvokeAsync<string, string>(
             "Jonas",
+            "1",
             TestContext.Current.CancellationToken
         );
 
@@ -93,7 +94,7 @@ public async Task LambdaHost_ServerInternalExceptions_AreCaughtAndReturnedAsErro
             .ContainSingle(ex =>
                 ex is InvalidOperationException
                 && ex.Message.Contains(
-                    "Unexpected request: GET http://http//localhost:3002/2018-06-01/runtime/invocation/next"
+                    "Unexpected request received from the Lambda HTTP handler: GET http://http//localhost:3002/2018-06-01/runtime/invocation/next"
                 )
             );
     }
@@ -163,8 +164,8 @@ await Assert.ThrowsAsync<TaskCanceledException>(() =>
     //         await using var factory = new LambdaApplicationFactory<Program>();
     //         await factory.TestServer.StartAsync(TestContext.Current.CancellationToken);
     //
-    //         var options = new LambdaClientOptions();
-    //         options.InvocationHeaderOptions.ClientWaitTimeout = TimeSpan.Zero;
+    //         var options = new LambdaServerOptions();
+    //         options.InvocationOptions.ClientWaitTimeout = TimeSpan.Zero;
     //
     //         await factory.TestServer.InvokeAsync<string, string>(
     //             "Jonas",
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServer.cs b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
index 46d96ad3..357598e8 100644
--- a/src/AwsLambda.Host.Testing/LambdaTestServer.cs
+++ b/src/AwsLambda.Host.Testing/LambdaTestServer.cs
@@ -16,7 +16,7 @@ public class LambdaTestServer : IAsyncDisposable
     /// <summary>
     /// Options used to configure how the server interacts with the Lambda.
     /// </summary>
-    private readonly LambdaClientOptions _clientOptions;
+    private readonly LambdaServerOptions _serverOptions;
 
     /// <summary>
     /// Task that represents the running Host application that has been captured.
@@ -102,7 +102,7 @@ internal LambdaTestServer(
         _state = ServerState.Created;
 
         _jsonSerializerOptions = DefaultLambdaJsonSerializerOptions.Create();
-        _clientOptions = new LambdaClientOptions();
+        _serverOptions = new LambdaServerOptions();
     }
 
     public IServiceProvider Services => _host!.Services;
@@ -142,10 +142,7 @@ public async Task<InitResponse> StartAsync(CancellationToken cancellationToken =
         if (_host is null)
             throw new InvalidOperationException("Host is not set.");
 
-        using var cts = CancellationTokenSource.CreateLinkedTokenSource(
-            cancellationToken,
-            _shutdownCts.Token
-        );
+        using var cts = LinkedCts(cancellationToken);
 
         _state = ServerState.Starting;
 
@@ -181,6 +178,7 @@ await TaskHelpers
 
     public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
         TEvent invokeEvent,
+        string? traceId = null,
         CancellationToken cancellationToken = default
     )
     {
@@ -189,27 +187,23 @@ public async Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
                 "TestServer is not Running and as such an event cannot be invoked."
             );
 
-        using var cts = CancellationTokenSource.CreateLinkedTokenSource(
-            cancellationToken,
-            _shutdownCts.Token
-        );
+        using var cts = LinkedCtsWithInvocationDeadline(cancellationToken);
+
+        traceId ??= Guid.NewGuid().ToString();
 
         // Generate unique request ID
         var requestId = GetRequestId();
 
         // Create the event response with Lambda headers
-        var eventResponse = CreateEventResponse(invokeEvent, requestId);
-        var deadlineUtc = DateTimeOffset.UtcNow.Add(
-            _clientOptions.InvocationHeaderOptions.ClientWaitTimeout
-        );
-
-        cts.CancelAfter(_clientOptions.InvocationHeaderOptions.ClientWaitTimeout);
+        var eventResponse = CreateEventResponse(invokeEvent, requestId, traceId);
+        var deadlineUtc = DateTimeOffset.UtcNow.Add(_serverOptions.FunctionTimeout);
 
         var pending = PendingInvocation.Create(requestId, eventResponse, deadlineUtc);
 
         _pendingInvocations.AddRequired(requestId, pending);
 
-        _pendingInvocationIds.Writer.TryWrite(requestId);
+        if (!_pendingInvocationIds.Writer.TryWrite(requestId))
+            throw new InvalidOperationException("Failed to enqueue pending invocation");
 
         var completion = await pending.ResponseTcs.Task.WaitAsync(cts.Token);
 
@@ -264,7 +258,7 @@ await TaskHelpers
     }
 
     //      ┌──────────────────────────────────────────────────────────┐
-    //      │                  Internal TestServer Logic                   │
+    //      │                Internal TestServer Logic                 │
     //      └──────────────────────────────────────────────────────────┘
 
     private async Task ProcessTransactionsAsync()
@@ -283,7 +277,7 @@ out var routeValues
                     )
                 )
                     throw new InvalidOperationException(
-                        $"Unexpected request: {transaction.Request.Method} {transaction.Request.RequestUri}"
+                        $"Unexpected request received from the Lambda HTTP handler: {transaction.Request.Method} {transaction.Request.RequestUri}"
                     );
 
                 switch (requestType!.Value)
@@ -382,7 +376,11 @@ private async Task HandlePostInitErrorAsync(LambdaHttpTransaction transaction)
         );
     }
 
-    private HttpResponseMessage CreateEventResponse<TEvent>(TEvent invokeEvent, string requestId)
+    private HttpResponseMessage CreateEventResponse<TEvent>(
+        TEvent invokeEvent,
+        string requestId,
+        string traceId
+    )
     {
         var response = new HttpResponseMessage(HttpStatusCode.OK)
         {
@@ -395,28 +393,20 @@ private HttpResponseMessage CreateEventResponse<TEvent>(TEvent invokeEvent, stri
         };
 
         // Add standard HTTP headers
-        response.Headers.Date = new DateTimeOffset(_clientOptions.InvocationHeaderOptions.Date);
-        response.Headers.TransferEncodingChunked = _clientOptions
-            .InvocationHeaderOptions
-            .TransferEncodingChunked;
+        response.Headers.Date = new DateTimeOffset(DateTime.UtcNow, TimeSpan.Zero);
+        response.Headers.TransferEncodingChunked = true;
 
         // Add custom Lambda runtime headers
         var deadlineMs = DateTimeOffset
-            .UtcNow.Add(_clientOptions.InvocationHeaderOptions.FunctionTimeout)
+            .UtcNow.Add(_serverOptions.FunctionTimeout)
             .ToUnixTimeMilliseconds();
         response.Headers.Add("Lambda-Runtime-Deadline-Ms", deadlineMs.ToString());
         response.Headers.Add("Lambda-Runtime-Aws-Request-Id", requestId);
-        response.Headers.Add(
-            "Lambda-Runtime-Trace-Id",
-            _clientOptions.InvocationHeaderOptions.TraceId
-        );
-        response.Headers.Add(
-            "Lambda-Runtime-Invoked-Function-Arn",
-            _clientOptions.InvocationHeaderOptions.FunctionArn
-        );
+        response.Headers.Add("Lambda-Runtime-Trace-Id", traceId);
+        response.Headers.Add("Lambda-Runtime-Invoked-Function-Arn", _serverOptions.FunctionArn);
 
         // Add any additional custom headers
-        foreach (var header in _clientOptions.InvocationHeaderOptions.AdditionalHeaders)
+        foreach (var header in _serverOptions.AdditionalHeaders)
             response.Headers.Add(header.Key, header.Value);
 
         return response;
@@ -468,4 +458,20 @@ private static HttpResponseMessage CreateSuccessResponse() =>
             ),
             Version = Version.Parse("1.1"),
         };
+
+    private CancellationTokenSource LinkedCtsWithInvocationDeadline(
+        CancellationToken cancellationTokens
+    )
+    {
+        var cts = CancellationTokenSource.CreateLinkedTokenSource(
+            cancellationTokens,
+            _shutdownCts.Token
+        );
+        cts.CancelAfter(_serverOptions.FunctionTimeout);
+
+        return cts;
+    }
+
+    private CancellationTokenSource LinkedCts(CancellationToken cancellationTokens) =>
+        CancellationTokenSource.CreateLinkedTokenSource(cancellationTokens, _shutdownCts.Token);
 }
diff --git a/src/AwsLambda.Host.Testing/LambdaTestServerExtensions.cs b/src/AwsLambda.Host.Testing/LambdaTestServerExtensions.cs
new file mode 100644
index 00000000..43220ecc
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/LambdaTestServerExtensions.cs
@@ -0,0 +1,16 @@
+namespace AwsLambda.Host.Testing;
+
+public static class LambdaTestServerExtensions
+{
+    extension(LambdaTestServer server)
+    {
+        public Task<InvocationResponse<TResponse>> InvokeAsync<TResponse, TEvent>(
+            TEvent invokeEvent,
+            CancellationToken cancellationToken = default
+        ) =>
+            server.InvokeAsync<TResponse, TEvent>(
+                invokeEvent,
+                cancellationToken: cancellationToken
+            );
+    }
+}
diff --git a/src/AwsLambda.Host.Testing/Options/LambdaClientOptions.cs b/src/AwsLambda.Host.Testing/Options/LambdaClientOptions.cs
deleted file mode 100644
index 025719b1..00000000
--- a/src/AwsLambda.Host.Testing/Options/LambdaClientOptions.cs
+++ /dev/null
@@ -1,12 +0,0 @@
-namespace AwsLambda.Host.Testing;
-
-/// <summary>
-/// Configuration options for the Lambda test client.
-/// </summary>
-public class LambdaClientOptions
-{
-    /// <summary>
-    /// Gets or sets the headers to include in Lambda invocation responses.
-    /// </summary>
-    public LambdaInvocationHeaderOptions InvocationHeaderOptions { get; set; } = new();
-}
diff --git a/src/AwsLambda.Host.Testing/Options/LambdaInvocationHeaderOptions.cs b/src/AwsLambda.Host.Testing/Options/LambdaInvocationOptions.cs
similarity index 54%
rename from src/AwsLambda.Host.Testing/Options/LambdaInvocationHeaderOptions.cs
rename to src/AwsLambda.Host.Testing/Options/LambdaInvocationOptions.cs
index 45379cc0..0c4e3ffd 100644
--- a/src/AwsLambda.Host.Testing/Options/LambdaInvocationHeaderOptions.cs
+++ b/src/AwsLambda.Host.Testing/Options/LambdaInvocationOptions.cs
@@ -3,7 +3,7 @@ namespace AwsLambda.Host.Testing;
 /// <summary>
 /// Headers returned in Lambda runtime API invocation responses.
 /// </summary>
-public class LambdaInvocationHeaderOptions
+public class LambdaInvocationOptions
 {
     /// <summary>
     /// Gets or sets additional custom headers to include in the Lambda invocation response.
@@ -11,12 +11,6 @@ public class LambdaInvocationHeaderOptions
     /// </summary>
     public Dictionary<string, string> AdditionalHeaders { get; set; } = new();
 
-    /// <summary>
-    /// Gets or sets the response date header.
-    /// Maps to the <c>Date</c> HTTP header. Defaults to current UTC time.
-    /// </summary>
-    public DateTime Date { get; set; } = DateTime.UtcNow;
-
     /// <summary>
     /// Gets or sets the ARN of the Lambda function being invoked.
     /// Maps to the <c>Lambda-Runtime-Invoked-Function-Arn</c> header.
@@ -32,32 +26,9 @@ public class LambdaInvocationHeaderOptions
     /// </summary>
     public TimeSpan FunctionTimeout { get; set; } = TimeSpan.FromMinutes(15);
 
-    /// <summary>
-    /// Gets or sets the maximum amount of time to wait for a response from the Lambda bootstrap.
-    /// Defaults to the function timeout plus a small buffer.
-    /// </summary>
-    public TimeSpan ClientWaitTimeout { get; set; } =
-        TimeSpan.FromMinutes(15).Add(TimeSpan.FromSeconds(5));
-
-    /// <summary>
-    /// Legacy property for client wait timeout. Prefer <see cref="ClientWaitTimeout"/>.
-    /// </summary>
-    [Obsolete("Use ClientWaitTimeout instead.")]
-    public TimeSpan InvocationTimeout
-    {
-        get => ClientWaitTimeout;
-        set => ClientWaitTimeout = value;
-    }
-
     /// <summary>
     /// Gets or sets the AWS X-Ray trace ID for distributed tracing.
     /// Maps to the <c>Lambda-Runtime-Trace-Id</c> header.
     /// </summary>
     public string TraceId { get; set; } = Guid.NewGuid().ToString();
-
-    /// <summary>
-    /// Gets or sets whether to use chunked transfer encoding.
-    /// Maps to the <c>Transfer-Encoding: chunked</c> HTTP header. Defaults to true.
-    /// </summary>
-    public bool TransferEncodingChunked { get; set; } = true;
 }
diff --git a/src/AwsLambda.Host.Testing/Options/LambdaServerOptions.cs b/src/AwsLambda.Host.Testing/Options/LambdaServerOptions.cs
new file mode 100644
index 00000000..2b92f9a2
--- /dev/null
+++ b/src/AwsLambda.Host.Testing/Options/LambdaServerOptions.cs
@@ -0,0 +1,28 @@
+namespace AwsLambda.Host.Testing;
+
+/// <summary>
+/// Configuration options for the Lambda test client.
+/// </summary>
+public class LambdaServerOptions
+{
+    /// <summary>
+    /// Gets or sets additional custom headers to include in the Lambda invocation response.
+    /// Use this to add any additional headers beyond the standard Lambda runtime headers.
+    /// </summary>
+    public Dictionary<string, string> AdditionalHeaders { get; set; } = new();
+
+    /// <summary>
+    /// Gets or sets the ARN of the Lambda function being invoked.
+    /// Maps to the <c>Lambda-Runtime-Invoked-Function-Arn</c> header.
+    /// </summary>
+    public string FunctionArn { get; set; } =
+        "arn:aws:lambda:us-west-2:123412341234:function:Function";
+
+    /// <summary>
+    /// Gets or sets the Lambda function timeout duration.
+    /// Maps to the <c>Lambda-Runtime-Deadline-Ms</c> header as Unix epoch milliseconds.
+    /// This determines when the Lambda function will timeout, calculated as the current time plus this duration.
+    /// Defaults to 15 minutes.
+    /// </summary>
+    public TimeSpan FunctionTimeout { get; set; } = TimeSpan.FromMinutes(15);
+}

From 404150faff55a066f82b5380b2fa6b9d1514083d Mon Sep 17 00:00:00 2001
From: Jonas Ha <jonas-ha@outlook.com>
Date: Wed, 10 Dec 2025 15:28:17 -0500
Subject: [PATCH 29/29] feat(testing): remove LambdaInvocationOptions class

- Deleted `LambdaInvocationOptions` class from `AwsLambda.Host.Testing`.
- Removed related properties like `AdditionalHeaders`, `FunctionArn`, `FunctionTimeout`, and `TraceId`.
- Streamlined testing code by eliminating unused infrastructure.
---
 .../Options/LambdaInvocationOptions.cs        | 34 -------------------
 1 file changed, 34 deletions(-)
 delete mode 100644 src/AwsLambda.Host.Testing/Options/LambdaInvocationOptions.cs

diff --git a/src/AwsLambda.Host.Testing/Options/LambdaInvocationOptions.cs b/src/AwsLambda.Host.Testing/Options/LambdaInvocationOptions.cs
deleted file mode 100644
index 0c4e3ffd..00000000
--- a/src/AwsLambda.Host.Testing/Options/LambdaInvocationOptions.cs
+++ /dev/null
@@ -1,34 +0,0 @@
-namespace AwsLambda.Host.Testing;
-
-/// <summary>
-/// Headers returned in Lambda runtime API invocation responses.
-/// </summary>
-public class LambdaInvocationOptions
-{
-    /// <summary>
-    /// Gets or sets additional custom headers to include in the Lambda invocation response.
-    /// Use this to add any additional headers beyond the standard Lambda runtime headers.
-    /// </summary>
-    public Dictionary<string, string> AdditionalHeaders { get; set; } = new();
-
-    /// <summary>
-    /// Gets or sets the ARN of the Lambda function being invoked.
-    /// Maps to the <c>Lambda-Runtime-Invoked-Function-Arn</c> header.
-    /// </summary>
-    public string FunctionArn { get; set; } =
-        "arn:aws:lambda:us-west-2:123412341234:function:Function";
-
-    /// <summary>
-    /// Gets or sets the Lambda function timeout duration.
-    /// Maps to the <c>Lambda-Runtime-Deadline-Ms</c> header as Unix epoch milliseconds.
-    /// This determines when the Lambda function will timeout, calculated as the current time plus this duration.
-    /// Defaults to 15 minutes.
-    /// </summary>
-    public TimeSpan FunctionTimeout { get; set; } = TimeSpan.FromMinutes(15);
-
-    /// <summary>
-    /// Gets or sets the AWS X-Ray trace ID for distributed tracing.
-    /// Maps to the <c>Lambda-Runtime-Trace-Id</c> header.
-    /// </summary>
-    public string TraceId { get; set; } = Guid.NewGuid().ToString();
-}
