ServiceBusSessionMessageActions not working with IsBatched=true for ServiceBusReceivedMessage

extensions/Worker.Extensions.ServiceBus/src/ServiceBusSessionMessageActionsConverter.cs

@@ -2,11 +2,12 @@
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
+using System.Text.Json;
 using System.Threading.Tasks;
 using Microsoft.Azure.Functions.Worker.Converters;
 using Microsoft.Azure.Functions.Worker.Extensions.Abstractions;
 using Microsoft.Azure.ServiceBus.Grpc;
-using System.Text.Json;
 
 namespace Microsoft.Azure.Functions.Worker
 {
@@ -15,7 +16,6 @@ namespace Microsoft.Azure.Functions.Worker
     /// </summary>
     [SupportsDeferredBinding]
     [SupportedTargetType(typeof(ServiceBusSessionMessageActions))]
-    [SupportedTargetType(typeof(ServiceBusSessionMessageActions[]))]
     internal class ServiceBusSessionMessageActionsConverter : IInputConverter
     {
         private readonly Settlement.SettlementClient _settlement;
@@ -29,11 +29,7 @@ public ValueTask<ConversionResult> ConvertAsync(ConverterContext context)
         {
             try
             {
-                var foundSessionId = context.FunctionContext.BindingContext.BindingData.TryGetValue("SessionId", out object? sessionId);
-                if (!foundSessionId)
-                {
-                    throw new InvalidOperationException($"Expecting SessionId within binding data and value was not present. Sessions must be enabled when binding to {nameof(ServiceBusSessionMessageActions)}.");
-                }
+                var sessionId = ParseSessionIdFromBindingData(context);
 
                 // Get the sessionLockedUntil property from the SessionActions binding data
                 var foundSessionActions = context.FunctionContext.BindingContext.BindingData.TryGetValue("SessionActions", out object? sessionActions);
@@ -49,7 +45,7 @@ public ValueTask<ConversionResult> ConvertAsync(ConverterContext context)
                     throw new InvalidOperationException("Expecting SessionLockedUntil within binding data of session actions and value was not present.");
                 }
 
-                var sessionActionResult = new ServiceBusSessionMessageActions(_settlement, sessionId!.ToString(), sessionLockedUntil.GetDateTimeOffset());
+                var sessionActionResult = new ServiceBusSessionMessageActions(_settlement, sessionId, sessionLockedUntil.GetDateTimeOffset());
                 var result = ConversionResult.Success(sessionActionResult);
                 return new ValueTask<ConversionResult>(result);
             }
@@ -58,5 +54,32 @@ public ValueTask<ConversionResult> ConvertAsync(ConverterContext context)
                 return new ValueTask<ConversionResult>(ConversionResult.Failed(exception));
             }
         }
+
+        private string ParseSessionIdFromBindingData(ConverterContext context)
+        {
+            // Try to resolve sessionId directly
+            var bindingData = context.FunctionContext.BindingContext.BindingData;
+            bindingData.TryGetValue("SessionId", out object? sessionId);
+
+            // If sessionId is not found and sessionIdRepeatedFieldArray has a value (isBatched = true), we can just parse the first sessionId from the array, as all the values are guaranteed to be the same.
+            // This is because there can be multiple messages but each message would belong to the same session.
+            // Note if web jobs extensions ever adds support for multiple sessions in a single batch, this logic will need to be updated.
+            if (sessionId == null && bindingData.TryGetValue("SessionIdArray", out object? sessionIdArray))
+            {
+                var sessionIdRepeatedArray = sessionIdArray as IList<string>;
+                if (sessionIdRepeatedArray is not null && sessionIdRepeatedArray.Count > 0)
+                {
+                    sessionId = sessionIdRepeatedArray[0]; // Use the first sessionId in the array
+                }
+            }
+
+            if (sessionId == null)
+            {
+                throw new InvalidOperationException(
+                    $"Expecting SessionId or SessionIdArray within binding data and value was not present. Sessions must be enabled when binding to {nameof(ServiceBusSessionMessageActions)}.");
+            }
+
+            return (string)sessionId;
+        }
     }
 }

test/Worker.Extensions.Tests/ServiceBus/ServiceBusMessageActionsTests.cs

@@ -103,7 +103,7 @@ public async Task PassingNullMessageThrows()
             await Assert.ThrowsAsync<ArgumentNullException>(async () => await messageActions.RenewMessageLockAsync(null));
         }
 
-        private class MockSettlementClient : Settlement.SettlementClient
+        internal class MockSettlementClient : Settlement.SettlementClient
         {
             private readonly string _lockToken;
             private readonly ByteString _propertiesToModify;

test/Worker.Extensions.Tests/ServiceBus/ServiceBusSessionMessageConverterTests.cs

@@ -0,0 +1,121 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using Microsoft.Azure.Functions.Worker.Converters;
+using static Microsoft.Azure.Functions.Worker.Extensions.Tests.ServiceBusMessageActionsTests;
+using Microsoft.Azure.Functions.Worker.Tests.Converters;
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Microsoft.Azure.Functions.Worker.Extensions.Tests
+{
+    public class ServiceBusSessionMessageConverterTests
+    {
+        internal sealed class TestBindingContext : BindingContext
+        {
+            public TestBindingContext(IReadOnlyDictionary<string, object?> input)
+            {
+                BindingData = input;
+            }
+
+            public override IReadOnlyDictionary<string, object?> BindingData { get; }
+        }
+
+        internal sealed class TestFunctionContext : FunctionContext
+        {
+            public TestFunctionContext(BindingContext bindingContext)
+            {
+                BindingContext = bindingContext;
+            }
+
+            public override BindingContext BindingContext { get; }
+
+            public override string InvocationId => throw new NotImplementedException();
+
+            public override string FunctionId => throw new NotImplementedException();
+
+            public override TraceContext TraceContext => throw new NotImplementedException();
+
+            public override RetryContext RetryContext => throw new NotImplementedException();
+
+            public override IServiceProvider InstanceServices { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
+
+            public override FunctionDefinition FunctionDefinition => throw new NotImplementedException();
+
+            public override IDictionary<object, object> Items { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
+
+            public override IInvocationFeatures Features => throw new NotImplementedException();
+        }
+
+        [Fact]
+        public async Task ConvertAsync_ReturnsSuccess()
+        {
+            var data = "{\"SessionLockedUntil\":\"2024-12-05T21:10:36.1193094+00:00\"}";
+
+            var bindingDataDictionary = new Dictionary<string, object>
+            {
+                { "SessionId", "test" },
+                { "SessionActions", JsonSerializer.Serialize(new
+                    {
+                        SessionLockedUntil = "2024-12-05T21:10:36.1193094+00:00"
+                    })
+                }
+            };
+            var context = new TestConverterContext(typeof(ServiceBusSessionMessageActions), data, new TestFunctionContext(new TestBindingContext(bindingDataDictionary)));
+            var converter = new ServiceBusSessionMessageActionsConverter(new MockSettlementClient("test"));
+            var result = await converter.ConvertAsync(context);
+
+            Assert.Equal(ConversionStatus.Succeeded, result.Status);
+            var output = result.Value as ServiceBusSessionMessageActions;
+            Assert.NotNull(output);
+        }
+
+        [Fact]
+        public async Task ConvertAsync_Batch_ForReceivedMessage_ReturnsSuccess()
+        {
+            var data = "{\"SessionLockedUntil\":\"2024-12-05T21:10:36.1193094+00:00\"}";
+
+            IList<string> repeatedField = new List<string> { "test" };
+
+            var bindingDataDictionary = new Dictionary<string, object>
+            {
+                { "SessionIdArray", repeatedField },
+                { "SessionActions", JsonSerializer.Serialize(new
+                    {
+                        SessionLockedUntil = "2024-12-05T21:10:36.1193094+00:00"
+                    })
+                }
+            };
+            var context = new TestConverterContext(typeof(ServiceBusSessionMessageActions), data, new TestFunctionContext(new TestBindingContext(bindingDataDictionary)));
+            var converter = new ServiceBusSessionMessageActionsConverter(new MockSettlementClient("test"));
+            var result = await converter.ConvertAsync(context);
+
+            Assert.Equal(ConversionStatus.Succeeded, result.Status);
+            var output = result.Value as ServiceBusSessionMessageActions;
+            Assert.NotNull(output);
+        }
+
+        [Fact]
+        public async Task ConvertAsync_ReturnsFailure_NoSessionId()
+        {
+            var data = "{\"SessionLockedUntil\":\"2024-12-05T21:10:36.1193094+00:00\"}";
+
+            var bindingDataDictionary = new Dictionary<string, object>
+            {
+                { "SessionActions", JsonSerializer.Serialize(new
+                    {
+                        SessionLockedUntil = "2024-12-05T21:10:36.1193094+00:00"
+                    })
+                }
+            };
+            var context = new TestConverterContext(typeof(ServiceBusSessionMessageActions), data, new TestFunctionContext(new TestBindingContext(bindingDataDictionary)));
+            var converter = new ServiceBusSessionMessageActionsConverter(new MockSettlementClient("test"));
+            var result = await converter.ConvertAsync(context);
+
+            Assert.Equal(ConversionStatus.Failed, result.Status);
+        }
+    }
+}