Add per-invocation async configure

eamsden · eamsden · commit 06d85f36ee3d · 2026-06-15T18:42:42.000-05:00
diff --git a/src/Temporalio.Extensions.Aws.Lambda/README.md b/src/Temporalio.Extensions.Aws.Lambda/README.md
@@ -33,9 +33,33 @@ public class Function
 }
 ```
 
-`configure` runs once when the delegate is created, which is normally during Lambda cold start. Each invocation creates a
-fresh Temporal client and worker, runs until the Lambda deadline minus `ShutdownDeadlineBuffer`, then shuts down and runs
-configured shutdown hooks.
+The synchronous `configure` callback shown above runs once when the delegate is created, which is normally during Lambda
+cold start. Use it for static worker setup such as task queues, workflow/activity registration, and options that can be
+shared across warm invocations. Each invocation creates a fresh Temporal client and worker, runs until the Lambda
+deadline minus `ShutdownDeadlineBuffer`, then shuts down and runs configured shutdown hooks.
+
+For setup that must be awaited per invocation, use the async overload:
+
+```csharp
+private static readonly Func<object?, ILambdaContext, Task> WorkerHandler =
+    TemporalLambdaWorker.CreateHandler(
+        new WorkerDeploymentVersion("payments-worker", "2026-05-27"),
+        async config =>
+        {
+            config.WorkerOptions.TaskQueue = "payments";
+            config.WorkerOptions.AddWorkflow<PaymentWorkflow>();
+            config.WorkerOptions.AddActivity(PaymentActivities.ChargeAsync);
+
+            config.ClientOptions.ApiKey = await LoadTemporalApiKeyAsync();
+            config.ShutdownHooks.Add(async cancellationToken =>
+            {
+                await FlushPerInvocationResourceAsync(cancellationToken);
+            });
+        });
+```
+
+The async `configure` callback is awaited once per Lambda invocation, before the Temporal client connects. It receives a
+fresh `LambdaWorkerConfig` each time, so use shutdown hooks for any per-invocation cleanup.
 
 ## Configuration
 
@@ -70,7 +94,8 @@ config.ClientOptions = new TemporalClientConnectOptions
 ```
 
 If `TEMPORAL_TASK_QUEUE` is present, it is used as the initial `WorkerOptions.TaskQueue`. You can still override the task
-queue in `configure`. If no task queue is set by the environment or by `configure`, handler creation fails.
+queue in `configure`. If no task queue is set by the environment or by `configure`, synchronous handler creation fails;
+with async configure, the invocation fails after the callback is awaited.
 
 `TemporalLambdaWorker.CreateHandler` requires a `WorkerDeploymentVersion` and always enables Worker Versioning by setting
 `WorkerOptions.DeploymentOptions` with `UseWorkerVersioning = true`. Use a deployment name and build ID that match your
@@ -80,7 +105,8 @@ value while enforcing the deployment version passed to `CreateHandler`.
 
 The helper applies Lambda-oriented worker defaults before `configure`, including lower concurrency, a 5 second graceful
 shutdown timeout, a smaller workflow cache, simple poller limits, and disabled eager activity execution. Values you set in
-`configure` override these defaults except for the enforced deployment version.
+`configure` override these defaults except for the enforced deployment version. With sync configure this happens once at
+handler creation; with async configure it happens for every invocation.
 
 ## Shutdown Hooks
 
@@ -94,7 +120,7 @@ config.ShutdownHooks.Add(async cancellationToken =>
 ```
 
 Hooks run in order after the worker has stopped. Hook failures are logged to the Lambda context logger and later hooks
-still run.
+still run. Hooks added from async configure are scoped to that invocation.
 
 ## Observability
 
diff --git a/src/Temporalio.Extensions.Aws.Lambda/TemporalLambdaWorker.cs b/src/Temporalio.Extensions.Aws.Lambda/TemporalLambdaWorker.cs
@@ -41,6 +41,23 @@ public static class TemporalLambdaWorker
                     LoadClientConnectOptions = options => LoadClientConnectOptions(options),
                 });
 
+        /// <summary>
+        /// Create an AWS Lambda handler that runs a Temporal worker for each invocation.
+        /// </summary>
+        /// <param name="version">Worker deployment version for this Lambda worker.</param>
+        /// <param name="configureAsync">Callback to configure client and worker options per invocation.</param>
+        /// <returns>A Lambda handler delegate.</returns>
+        public static Func<object?, ILambdaContext, Task> CreateHandler(
+            WorkerDeploymentVersion version,
+            Func<LambdaWorkerConfig, Task> configureAsync) =>
+            CreateHandler(
+                version,
+                configureAsync,
+                new TemporalLambdaWorkerHandlerOptions
+                {
+                    LoadClientConnectOptions = options => LoadClientConnectOptions(options),
+                });
+
         /// <summary>
         /// Load Temporal client connection options using AWS Lambda-aware config file resolution.
         /// </summary>
@@ -80,14 +97,44 @@ public static TemporalClientConnectOptions LoadClientConnectOptions(
             WorkerDeploymentVersion version,
             Action<LambdaWorkerConfig> configure,
             TemporalLambdaWorkerHandlerOptions handlerOptions)
+        {
+            ValidateCreateHandlerArgs(version, configure, nameof(configure), handlerOptions);
+            var config = CreateConfig(version, handlerOptions);
+            configure(config);
+            var state = PrepareHandlerState(version, config, handlerOptions);
+            return state.HandleAsync;
+        }
+
+        /// <summary>
+        /// Create an AWS Lambda handler with overridable internals for tests.
+        /// </summary>
+        /// <param name="version">Worker deployment version for this Lambda worker.</param>
+        /// <param name="configureAsync">Callback to configure client and worker options per invocation.</param>
+        /// <param name="handlerOptions">Internal handler options.</param>
+        /// <returns>A Lambda handler delegate.</returns>
+        internal static Func<object?, ILambdaContext, Task> CreateHandler(
+            WorkerDeploymentVersion version,
+            Func<LambdaWorkerConfig, Task> configureAsync,
+            TemporalLambdaWorkerHandlerOptions handlerOptions)
+        {
+            ValidateCreateHandlerArgs(version, configureAsync, nameof(configureAsync), handlerOptions);
+            var state = new AsyncLambdaWorkerHandlerState(version, configureAsync, handlerOptions);
+            return state.HandleAsync;
+        }
+
+        private static void ValidateCreateHandlerArgs(
+            WorkerDeploymentVersion version,
+            object configure,
+            string configureParamName,
+            TemporalLambdaWorkerHandlerOptions handlerOptions)
         {
             if (version == null)
             {
                 throw new ArgumentNullException(nameof(version));
             }
             if (configure == null)
             {
-                throw new ArgumentNullException(nameof(configure));
+                throw new ArgumentNullException(configureParamName);
             }
             if (handlerOptions == null)
             {
@@ -101,7 +148,12 @@ public static TemporalClientConnectOptions LoadClientConnectOptions(
             {
                 throw new ArgumentException("Build ID must be set", nameof(version));
             }
+        }
 
+        private static LambdaWorkerConfig CreateConfig(
+            WorkerDeploymentVersion version,
+            TemporalLambdaWorkerHandlerOptions handlerOptions)
+        {
             var loadClientConnectOptions = handlerOptions.LoadClientConnectOptions;
             var config = new LambdaWorkerConfig(
                 loadClientConnectOptions == null ?
@@ -114,9 +166,7 @@ public static TemporalClientConnectOptions LoadClientConnectOptions(
             }
             ApplyDeploymentVersion(config.WorkerOptions, version);
 
-            configure(config);
-            var state = PrepareHandlerState(version, config, handlerOptions);
-            return state.HandleAsync;
+            return config;
         }
 
         private static LambdaWorkerHandlerState PrepareHandlerState(
@@ -345,5 +395,35 @@ private async Task RunShutdownHooksAsync(ILambdaContext context)
                 }
             }
         }
+
+        private sealed class AsyncLambdaWorkerHandlerState
+        {
+            private readonly WorkerDeploymentVersion version;
+            private readonly Func<LambdaWorkerConfig, Task> configureAsync;
+            private readonly TemporalLambdaWorkerHandlerOptions handlerOptions;
+
+            public AsyncLambdaWorkerHandlerState(
+                WorkerDeploymentVersion version,
+                Func<LambdaWorkerConfig, Task> configureAsync,
+                TemporalLambdaWorkerHandlerOptions handlerOptions)
+            {
+                this.version = version;
+                this.configureAsync = configureAsync;
+                this.handlerOptions = handlerOptions;
+            }
+
+            public async Task HandleAsync(object? input, ILambdaContext context)
+            {
+                if (context == null)
+                {
+                    throw new ArgumentNullException(nameof(context));
+                }
+
+                var config = CreateConfig(version, handlerOptions);
+                await configureAsync(config).ConfigureAwait(false);
+                var state = PrepareHandlerState(version, config, handlerOptions);
+                await state.HandleAsync(input, context).ConfigureAwait(false);
+            }
+        }
     }
 }
diff --git a/tests/Temporalio.Tests/Extensions/Aws/Lambda/TemporalLambdaWorkerTests.cs b/tests/Temporalio.Tests/Extensions/Aws/Lambda/TemporalLambdaWorkerTests.cs
@@ -752,6 +752,144 @@ public async Task Invoke_ShutdownHooksRunInOrderPerInvocationAndContinueAfterFai
         Assert.Equal(2, workerCreations);
     }
 
+    [Fact]
+    public async Task CreateHandler_AsyncConfigureRunsPerInvocationWithFreshConfig()
+    {
+        var configureCalls = 0;
+        var capturedConfigs = new List<LambdaWorkerConfig>();
+        var capturedTargets = new List<string?>();
+        var capturedTaskQueues = new List<string?>();
+        var hookCalls = new List<string>();
+        var handler = TemporalLambdaWorker.CreateHandler(
+            Version,
+            async config =>
+            {
+                await Task.Yield();
+                var call = ++configureCalls;
+                capturedConfigs.Add(config);
+                Assert.Equal("env-task-queue", config.WorkerOptions.TaskQueue);
+
+                config.ClientOptions.TargetHost = $"target-{call}";
+                config.WorkerOptions.TaskQueue = $"task-queue-{call}";
+                config.ShutdownHooks.Add(_ =>
+                {
+                    hookCalls.Add($"hook-{call}");
+                    return Task.CompletedTask;
+                });
+            },
+            new TemporalLambdaWorkerHandlerOptions
+            {
+                GetEnvironmentVariable = name =>
+                    name == "TEMPORAL_TASK_QUEUE" ? "env-task-queue" : null,
+                ConnectClientAsync = options =>
+                {
+                    capturedTargets.Add(options.TargetHost);
+                    return Task.FromResult<object>(new object());
+                },
+                CreateWorker = (_, options) =>
+                {
+                    capturedTaskQueues.Add(options.TaskQueue);
+                    return new FakeLambdaWorker(_ => Task.CompletedTask);
+                },
+            });
+
+        await handler(null, new FakeLambdaContext());
+        await handler(null, new FakeLambdaContext());
+
+        Assert.Equal(2, configureCalls);
+        Assert.Equal(2, capturedConfigs.Count);
+        Assert.NotSame(capturedConfigs[0], capturedConfigs[1]);
+        Assert.Equal(new[] { "target-1", "target-2" }, capturedTargets);
+        Assert.Equal(new[] { "task-queue-1", "task-queue-2" }, capturedTaskQueues);
+        Assert.Equal(new[] { "hook-1", "hook-2" }, hookCalls);
+    }
+
+    [Fact]
+    public async Task CreateHandler_AsyncConfigureErrorsSurfaceOnInvocation()
+    {
+        var configureCalls = 0;
+        var handler = TemporalLambdaWorker.CreateHandler(
+            Version,
+            async config =>
+            {
+                _ = config;
+                await Task.Yield();
+                configureCalls++;
+                throw new InvalidOperationException("bad config");
+            },
+            new TemporalLambdaWorkerHandlerOptions());
+
+        var error = await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            handler(null, new FakeLambdaContext()));
+        Assert.Equal("bad config", error.Message);
+        Assert.Equal(1, configureCalls);
+    }
+
+    [Fact]
+    public async Task CreateHandler_AsyncConfigureValidatesTaskQueueOnInvocation()
+    {
+        var handler = TemporalLambdaWorker.CreateHandler(
+            Version,
+            _ => Task.CompletedTask,
+            new TemporalLambdaWorkerHandlerOptions
+            {
+                GetEnvironmentVariable = _ => null,
+            });
+
+        await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            handler(null, new FakeLambdaContext()));
+    }
+
+    [Fact]
+    public async Task CreateHandler_AsyncConfigureShutdownHooksRunAfterFailures()
+    {
+        var hookCalls = new List<string>();
+        var connectFailureHandler = TemporalLambdaWorker.CreateHandler(
+            Version,
+            config =>
+            {
+                config.WorkerOptions.TaskQueue = "task-queue";
+                config.ShutdownHooks.Add(_ =>
+                {
+                    hookCalls.Add("connect");
+                    return Task.CompletedTask;
+                });
+                return Task.CompletedTask;
+            },
+            new TemporalLambdaWorkerHandlerOptions
+            {
+                ConnectClientAsync = _ =>
+                    throw new InvalidOperationException("connect failed"),
+            });
+
+        await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            connectFailureHandler(null, new FakeLambdaContext()));
+
+        var workerFailureHandler = TemporalLambdaWorker.CreateHandler(
+            Version,
+            config =>
+            {
+                config.WorkerOptions.TaskQueue = "task-queue";
+                config.ShutdownHooks.Add(_ =>
+                {
+                    hookCalls.Add("worker");
+                    return Task.CompletedTask;
+                });
+                return Task.CompletedTask;
+            },
+            new TemporalLambdaWorkerHandlerOptions
+            {
+                ConnectClientAsync = _ => Task.FromResult<object>(new object()),
+                CreateWorker = (_, _) => new FakeLambdaWorker(
+                    _ => throw new InvalidOperationException("worker failed")),
+            });
+
+        await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            workerFailureHandler(null, new FakeLambdaContext()));
+
+        Assert.Equal(new[] { "connect", "worker" }, hookCalls);
+    }
+
     private static Func<object?, ILambdaContext, Task> CreateCapturingHandler(
         Action<LambdaWorkerConfig> configure,
         Action<TemporalClientConnectOptions> captureClientOptions) =>
