From 7dd1eb90ad34e486a46432880144d1139269d86b Mon Sep 17 00:00:00 2001
From: Jon Sequeira <jonsequeira@gmail.com>
Date: Mon, 13 Feb 2023 19:29:05 -0800
Subject: [PATCH] fix handling of completions that rely on nested commands
 (#2727)

* fix handling of completions that rely on nested commands

* add logic to compensate for #2728
---
 NotebookTestScript.dib                        | 15 +++++++--
 .../KernelExtensions.cs                       | 33 +++++++++++++------
 .../src/interactiveClient.ts                  |  2 +-
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/NotebookTestScript.dib b/NotebookTestScript.dib
index 82191e3c85..c4d39d5ac7 100644
--- a/NotebookTestScript.dib
+++ b/NotebookTestScript.dib
@@ -1,6 +1,6 @@
 #!meta
 
-{"kernelInfo":{"defaultKernelName":null,"items":[{"name":"csharp","languageName":"C#","aliases":["c#","cs"]},{"name":"fsharp","languageName":"F#","aliases":["f#","fs"]},{"name":"pwsh","languageName":"PowerShell","aliases":["powershell"]},{"name":"javascript","languageName":"JavaScript","aliases":["js"]},{"name":"html","languageName":"HTML"},{"name":"sql","languageName":"SQL"},{"name":"kql","languageName":"KQL"},{"name":"mermaid","languageName":"Mermaid"},{"name":"http","languageName":"HTTP"},{"name":"value"}]}}
+{"kernelInfo":{"defaultKernelName":"csharp","items":[{"name":"csharp","aliases":[]}]}}
 
 #!markdown
 
@@ -188,7 +188,7 @@ Console.Error.Write("\x1b[0m");
 #!markdown
 
 # vscode kernel and user input prompts
-This cell should prompt the user for input
+This cell should prompt the user for input. Type something into the input field. Whatever you typed should be displayed in the output.
 
 #!csharp
 
@@ -221,7 +221,16 @@ value_from_input
 
 #!markdown
 
-Now check the variable explorer. Is the value there in the `#!value` kernel?
+Now check the variable explorer. Is the value there in the `value` kernel?
+
+#!markdown
+
+This should prompt you for input. Type something into the input field. Whatever you typed should be displayed in the output.
+
+#!csharp
+
+#!set --name csharpVarFromInput --value @input:"Please enter a value"
+csharpVarFromInput
 
 #!markdown
 
diff --git a/src/Microsoft.DotNet.Interactive/KernelExtensions.cs b/src/Microsoft.DotNet.Interactive/KernelExtensions.cs
index d41726fe49..4868de002f 100644
--- a/src/Microsoft.DotNet.Interactive/KernelExtensions.cs
+++ b/src/Microsoft.DotNet.Interactive/KernelExtensions.cs
@@ -283,25 +283,38 @@ private static void ConfigureAndAddSetMagicCommand<T>(T kernel) where T : Kernel
 
         valueOption.AddCompletions(_ =>
         {
-            // FIX: (ConfigureAndAddSetMagicCommand) why is this getting invoked this option wasn't specified?
-
             if (kernel.ParentKernel is { } composite)
             {
                 var getValueTasks = composite.ChildKernels
                                              .Where(
                                                  k => k != kernel &&
                                                       k.KernelInfo.SupportsCommand(nameof(RequestValueInfos)))
-                                             .Select(async k => await k.SendAsync(new RequestValueInfos()));
+                                             .Select(async k => await k.SendAsync(new RequestValueInfos(k.Name)));
 
                 var tasks = Task.WhenAll(getValueTasks).GetAwaiter().GetResult();
 
-                return tasks
-                       .Select(t => t.Events.OfType<ValueInfosProduced>())
-                       .SelectMany(events => events.Select(e => new { e.Command.TargetKernelName, e.ValueInfos }))
-                       .SelectMany(x => x.ValueInfos.Select(i => $"@{x.TargetKernelName}:{i.Name}"))
-                       .OrderBy(x => x)
-                       .Select(n => new CompletionItem(n))
-                       .ToArray();
+                var x = tasks
+                        .Select(t => t.Events.OfType<ValueInfosProduced>())
+                        .SelectMany(events => events.Select(e => new { e.Command, e.ValueInfos }))
+                        .SelectMany(x =>
+                        {
+                            var kernelName = x.Command.TargetKernelName;
+
+                            // TODO: (ConfigureAndAddSetMagicCommand) this is compensating for https://github.com/dotnet/interactive/issues/2728
+                            if (kernelName is null &&
+                                kernel.RootKernel is CompositeKernel root &&
+                                root.ChildKernels.TryGetByUri(x.Command.DestinationUri, out var k))
+                            {
+                                kernelName = k.Name;
+                            }
+
+                            return x.ValueInfos.Select(i => $"@{kernelName}:{i.Name}");
+                        })
+                        .OrderBy(x => x)
+                        .Select(n => new CompletionItem(n))
+                        .ToArray();
+
+                return x;
             }
 
             return Array.Empty<CompletionItem>();
diff --git a/src/polyglot-notebooks-vscode-common/src/interactiveClient.ts b/src/polyglot-notebooks-vscode-common/src/interactiveClient.ts
index 1ac9dee036..089cc308f1 100644
--- a/src/polyglot-notebooks-vscode-common/src/interactiveClient.ts
+++ b/src/polyglot-notebooks-vscode-common/src/interactiveClient.ts
@@ -333,7 +333,7 @@ export class InteractiveClient {
             token = token || this.getNextToken();
             const id = Guid.create().toString();
             let disposable = this.subscribeToKernelTokenEvents(token, eventEnvelope => {
-                if (eventEnvelope.command?.token === token) {
+                if (eventEnvelope.command?.token === token && eventEnvelope.eventType === expectedEventType) {
                     switch (eventEnvelope.eventType) {
                         case CommandFailedType:
                             if (!handled) {