Refactor NL2Fx Handling in Language Server to support more dynamic NL Handler Creation based on IPowerFxScope (#2134)

This pr performs necessary wiring to enable creating NL Handler for
Nl2Fx and FX2NL based on the current IPowerFxScope language server is
operating on. This is necessary to support Nl2FX scenarios in Power Apps
and the current approach which uses Nl2FxImplementation property is not
a thread safe. With the plan to perform Nl2Fx operations in parallel and
given that we only create 1 instance of Language Server SDK, two
concurrent Nl2FX tasks can overwrite Nl2FxImplementation unexpectedly.
The new approach in this pr scopes the creation of handler to every
Nl2FX task. It allows the creation of handler based on the scope itself
which has "host" specific entities which are required to create a more
meaningful NlHandler. For Power Apps, "host" contains document, control
and property which are critical to context building for Nl2Fx model
inside NL Handler.

src/libraries/Microsoft.PowerFx.LanguageServerProtocol/LanguageServer/INLHandlerFactory.cs

@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using Microsoft.PowerFx.Intellisense;
+using Microsoft.PowerFx.LanguageServerProtocol.Protocol;
+
+namespace Microsoft.PowerFx.LanguageServerProtocol
+{
+    /// <summary>
+    /// Represents a factory for creating NLHandler.
+    /// </summary>
+    public interface INLHandlerFactory
+    {
+        /// <summary>
+        /// Get the NLHandler from the given scope.
+        /// </summary>
+        /// <param name="scope">IPowerFxScope instance.</param>
+        /// <param name="nlParams">NL operation params.</param>
+        /// <returns>NLHandler for the given scope.</returns>
+        NLHandler GetNLHandler(IPowerFxScope scope, BaseNLParams nlParams);
+    }
+}

src/libraries/Microsoft.PowerFx.LanguageServerProtocol/LanguageServer/LanguageServer.cs

@@ -35,6 +35,11 @@ namespace Microsoft.PowerFx.LanguageServerProtocol
     {
         private const char EOL = '\n';
 
+        /// <summary>
+        /// This const represents the dummy formula that is used to create an infrastructure needed to get the symbols for Nl2Fx operation.
+        /// </summary>
+        internal static readonly string Nl2FxDummyFormula = "\"f7979178-07f0-424d-8f8b-00fee6fd19b8\"";
+
         public delegate void SendToClient(string data);
 
         private readonly SendToClient _sendToClient;
@@ -60,9 +65,15 @@ namespace Microsoft.PowerFx.LanguageServerProtocol
 
         /// <summary>
         /// If set, provides the handler for $/nlSuggestion message.
+        /// Note: This is not a thread safe. Consider using the NlHandlerFactory.
         /// </summary>
         public NLHandler NL2FxImplementation { get; set; }
 
+        /// <summary>
+        /// A factory to get the NLHandler from the given scope.
+        /// </summary>
+        public INLHandlerFactory NLHandlerFactory { get; init; }
+
         public LanguageServer(SendToClient sendToClient, IPowerFxScopeFactory scopeFactory, Action<string> logger = null)
         {
             Contracts.AssertValue(sendToClient);
@@ -464,7 +475,7 @@ namespace Microsoft.PowerFx.LanguageServerProtocol
             {
                 var result = new CustomGetCapabilitiesResult();
 
-                var nl = _parent.NL2FxImplementation;
+                var nl = GetNLHandler(_parent, scope, null);
                 if (nl != null)
                 {
                     result.SupportsNL2Fx = nl.SupportsNL2Fx;
@@ -484,13 +495,13 @@ namespace Microsoft.PowerFx.LanguageServerProtocol
 
             protected override CustomNL2FxResult Handle(IPowerFxScope scope, CustomNL2FxParams request)
             {
-                var nl = _parent.NL2FxImplementation;
+                var nl = GetNLHandler(_parent, scope, request);
                 if (nl == null || !nl.SupportsNL2Fx)
                 {
                     throw new NotSupportedException($"NL2Fx not enabled");
                 }
 
-                var check = scope.Check("1"); // just need to get the symbols 
+                var check = scope.Check(Nl2FxDummyFormula); // just need to get the symbols 
                 var summary = check.ApplyGetContextSummary();
 
                 var req = new NL2FxParameters
@@ -501,7 +512,7 @@ namespace Microsoft.PowerFx.LanguageServerProtocol
                 };
 
                 CancellationToken cancel = default;
-                var result = _parent.NL2FxImplementation.NL2FxAsync(req, cancel)
+                var result = nl.NL2FxAsync(req, cancel)
                     .ConfigureAwait(false).GetAwaiter().GetResult();
 
                 FinalCheck(scope, result);
@@ -537,7 +548,7 @@ namespace Microsoft.PowerFx.LanguageServerProtocol
 
             protected override CustomFx2NLResult Handle(IPowerFxScope scope, CustomFx2NLParams request)
             {
-                var nl = _parent.NL2FxImplementation;
+                var nl = GetNLHandler(_parent, scope, request);
                 if (nl == null || !nl.SupportsFx2NL)
                 {
                     throw new NotSupportedException($"NL2Fx not enabled");
@@ -552,6 +563,11 @@ namespace Microsoft.PowerFx.LanguageServerProtocol
             }
         }
 
+        private static NLHandler GetNLHandler(LanguageServer server, IPowerFxScope scope, BaseNLParams nlParams)
+        {
+            return server.NLHandlerFactory?.GetNLHandler(scope, nlParams) ?? server.NL2FxImplementation;
+        }
+
         private void HandleCodeActionRequest(string id, string paramsJson)
         {
             _logger?.Invoke($"[PFX] HandleCodeActionRequest: id={id ?? "<null>"}, paramsJson={paramsJson ?? "<null>"}");

src/libraries/Microsoft.PowerFx.LanguageServerProtocol/Protocol/BaseNLParams.cs

@@ -0,0 +1,13 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+namespace Microsoft.PowerFx.LanguageServerProtocol.Protocol
+{
+    public class BaseNLParams
+    {
+        /// <summary>
+        /// Additional context for NL operation. Usually, a stringified JSON object.
+        /// </summary>
+        public string Context { get; set; } = null;
+    }
+}

src/libraries/Microsoft.PowerFx.LanguageServerProtocol/Protocol/CustomFx2NLParams.cs

@@ -10,7 +10,7 @@ namespace Microsoft.PowerFx.LanguageServerProtocol.Protocol
     /// <summary>
     /// Incoming LSP payload for a NL request. See <see cref="CustomProtocolNames.FX2NL"/>.
     /// </summary>
-    public class CustomFx2NLParams : IHasTextDocument
+    public class CustomFx2NLParams : BaseNLParams, IHasTextDocument
     {
         /// <summary>
         /// The document that was opened. Just need Uri. 

src/libraries/Microsoft.PowerFx.LanguageServerProtocol/Protocol/CustomNL2FxParams.cs

@@ -10,7 +10,7 @@ namespace Microsoft.PowerFx.LanguageServerProtocol.Protocol
     /// <summary>
     /// Incoming LSP payload for a NL request. See <see cref="CustomProtocolNames.NL2FX"/>.
     /// </summary>
-    public class CustomNL2FxParams : IHasTextDocument
+    public class CustomNL2FxParams : BaseNLParams, IHasTextDocument
     {
         /// <summary>
         /// The document that was opened. Just need Uri. 

src/tests/Microsoft.PowerFx.Interpreter.Tests/LanguageServiceProtocol/LanguageServerTests.cs

@@ -1492,11 +1492,7 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol.Tests
                 testNLHandler = null;
             }
 
-            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory)
-            {
-                NL2FxImplementation = testNLHandler
-            };
-
+            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory, new TestNlHandlerFactory(testNLHandler));
             List<Exception> exList = new List<Exception>();
             testServer.LogUnhandledExceptionHandler += (Exception ex) => exList.Add(ex);
 
@@ -1646,7 +1642,7 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol.Tests
             });
         }
 
-        private static string FX2NlMessageJson(string documentUri)
+        private static string FX2NlMessageJson(string documentUri, string context = null)
         {
             return JsonSerializer.Serialize(new
             {
@@ -1661,15 +1657,39 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol.Tests
                         LanguageId = "powerfx",
                         Version = 1
                     },
-                    Expression = "Score > 3"
+                    Expression = "Score > 3",
+                    Context = context
                 }
             });
         }
 
+        private class TestNlHandlerFactory : INLHandlerFactory
+        {
+            private readonly TestNLHandler _handler;
+
+            public IPowerFxScope Scope { get; private set; }
+
+            public BaseNLParams NLParams { get; private set; }
+
+            public TestNlHandlerFactory(TestNLHandler handler)
+            {
+                _handler = handler;
+            }
+
+            public NLHandler GetNLHandler(IPowerFxScope scope, BaseNLParams nlParams)
+            {
+                Scope = scope;
+                NLParams = nlParams;
+                return _handler;
+            }
+        }
+
         [Theory]
-        [InlineData("Score < 50", true)]
-        [InlineData("missing < 50", false)] // doesn't compile, should get filtered out by LSP 
-        public void TestNL2FX(string expectedExpr, bool success)
+        [InlineData("Score < 50", true, true)]
+        [InlineData("missing < 50", false, true)] // doesn't compile, should get filtered out by LSP 
+        [InlineData("Score < 50", true, false)]
+        [InlineData("missing < 50", false, false)]
+        public void TestNL2FX(string expectedExpr, bool success, bool useFactory)
         {
             var documentUri = "powerfx://app?context=1";
 
@@ -1685,7 +1705,7 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol.Tests
             var scopeFactory = new TestPowerFxScopeFactory((string documentUri) => dict[documentUri]);
 
             var testNLHandler = new TestNLHandler { Expected = expectedExpr };
-            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory)
+            var testServer = useFactory ? new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory, new TestNlHandlerFactory(testNLHandler)) : new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory)
             {
                 NL2FxImplementation = testNLHandler
             };
@@ -1766,10 +1786,7 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol.Tests
             var scopeFactory = new TestPowerFxScopeFactory((string documentUri) => dict[documentUri]);
 
             var testNLHandler = new TestNLHandler { Throw = true }; // simulate error
-            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory)
-            {
-                NL2FxImplementation = testNLHandler
-            };
+            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory, new TestNlHandlerFactory(testNLHandler));
 
             List<Exception> exList = new List<Exception>();
             testServer.LogUnhandledExceptionHandler += (Exception ex) => exList.Add(ex);
@@ -1783,6 +1800,49 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol.Tests
             Assert.NotNull(errorResponse.Error);
         }
 
+        [Fact]
+        public void TestNlHandlerCreation()
+        {
+            var documentUri = "powerfx://app?context=1";
+            var expectedExpr = "sentence";
+
+            var engine = new Engine();
+            var symbols = new SymbolTable();
+            symbols.AddVariable("Score", FormulaType.Number);
+            var editor = engine.CreateEditorScope(symbols: symbols);
+
+            var dict = new Dictionary<string, EditorContextScope>
+            {
+                { documentUri, editor }
+            };
+            var scopeFactory = new TestPowerFxScopeFactory((string documentUri) => dict[documentUri]);
+           
+            var testNLHandler = new TestNLHandler { Expected = expectedExpr };
+            var nlHandlerFactory = new TestNlHandlerFactory(testNLHandler);
+
+            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory, nlHandlerFactory);
+
+            List<Exception> exList = new List<Exception>();
+            testServer.LogUnhandledExceptionHandler += (Exception ex) => exList.Add(ex);
+
+            testServer.OnDataReceived(FX2NlMessageJson(documentUri, JsonSerializer.Serialize(new
+            {
+                Age = 24,
+                Name = "Foobar"
+            })));
+
+            var nlContext = nlHandlerFactory.NLParams.Context;
+            Assert.NotNull(nlContext);
+            Assert.NotNull(nlHandlerFactory.Scope);
+
+            var document = JsonDocument.Parse(nlContext);
+            var root = document.RootElement;
+            Assert.True(root.TryGetProperty("Age", out var age));
+            Assert.Equal(24, age.GetInt32());
+            Assert.True(root.TryGetProperty("Name", out var name));
+            Assert.Equal("Foobar", name.GetString());
+        }
+
         [Fact]
         public void TestFx2NL()
         {
@@ -1801,10 +1861,7 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol.Tests
             var scopeFactory = new TestPowerFxScopeFactory((string documentUri) => dict[documentUri]);
 
             var testNLHandler = new TestNLHandler { Expected = expectedExpr };
-            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory)
-            {
-                NL2FxImplementation = testNLHandler
-            };
+            var testServer = new TestLanguageServer(_output, _sendToClientData.Add, scopeFactory, new TestNlHandlerFactory(testNLHandler));
 
             List<Exception> exList = new List<Exception>();
             testServer.LogUnhandledExceptionHandler += (Exception ex) => exList.Add(ex);

src/tests/Microsoft.PowerFx.Interpreter.Tests/LanguageServiceProtocol/PublicSurfaceTests.cs

@@ -36,6 +36,7 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol
 
                 "Microsoft.PowerFx.LanguageServerProtocol.NLHandler",
                 "Microsoft.PowerFx.LanguageServerProtocol.NL2FxParameters",
+                "Microsoft.PowerFx.LanguageServerProtocol.INLHandlerFactory",
 
                 // Internal
                 "Microsoft.PowerFx.LanguageServerProtocol.JsonRpcHelper",
@@ -59,6 +60,7 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol
                 "Microsoft.PowerFx.LanguageServerProtocol.Protocol.CustomProtocolNames",
                 "Microsoft.PowerFx.LanguageServerProtocol.Protocol.CustomGetCapabilitiesParams",
                 "Microsoft.PowerFx.LanguageServerProtocol.Protocol.CustomGetCapabilitiesResult",
+                "Microsoft.PowerFx.LanguageServerProtocol.Protocol.BaseNLParams",
                 "Microsoft.PowerFx.LanguageServerProtocol.Protocol.CustomFx2NLParams",
                 "Microsoft.PowerFx.LanguageServerProtocol.Protocol.CustomFx2NLResult",
                 "Microsoft.PowerFx.LanguageServerProtocol.Protocol.CustomNL2FxParams",

src/tests/Microsoft.PowerFx.Interpreter.Tests/LanguageServiceProtocol/TestLanguageServer.cs

@@ -9,9 +9,10 @@ namespace Microsoft.PowerFx.Tests.LanguageServiceProtocol
 {
     public class TestLanguageServer : LanguageServer
     {
-        public TestLanguageServer(ITestOutputHelper output, SendToClient sendToClient, IPowerFxScopeFactory scopeFactory)
+        public TestLanguageServer(ITestOutputHelper output, SendToClient sendToClient, IPowerFxScopeFactory scopeFactory, INLHandlerFactory nlHandlerFactory = null)
             : base(sendToClient, scopeFactory, (string s) => output.WriteLine(s))
         {            
+            NLHandlerFactory = nlHandlerFactory;
         }
 
         public int TestGetCharPosition(string expression, int position) => GetCharPosition(expression, position);