bubble events for specific commands and their descendants when proxied

Author: jonsequitur

src/Microsoft.DotNet.Interactive.Tests/DirectiveTests.cs

@@ -8,6 +8,7 @@
 using System.Linq;
 using System.Threading.Tasks;
 using FluentAssertions;
+using Microsoft.DotNet.Interactive.Commands;
 using Microsoft.DotNet.Interactive.CSharp;
 using Microsoft.DotNet.Interactive.Events;
 using Microsoft.DotNet.Interactive.Formatting;
@@ -313,4 +314,29 @@ public async Task New_directives_can_be_added_after_older_ones_have_been_evaluat
         oneWasCalled.Should().BeTrue();
         twoWasCalled.Should().BeTrue();
     }
+
+    [Fact]
+    public async Task Results_of_commands_sent_within_directive_handlers_contain_expected_events()
+    {
+        var cSharpKernel = new CSharpKernel();
+
+        using var kernel = new CompositeKernel { cSharpKernel };
+
+        var directive = new Command("#!dostuff");
+
+        ReturnValueProduced returnValueProduced = null;
+
+        directive.SetHandler(async () =>
+        {
+            var result = await cSharpKernel.SendAsync(new SubmitCode("123"));
+
+            returnValueProduced = result.Events.OfType<ReturnValueProduced>().SingleOrDefault();
+        });
+
+        kernel.AddDirective(directive);
+
+        await kernel.SendAsync(new SubmitCode("#!dostuff"));
+
+        returnValueProduced.Should().NotBeNull();
+    }
 }

src/Microsoft.DotNet.Interactive.Tests/KernelCommandNestingTests.cs

@@ -124,6 +124,33 @@ public async Task Commands_sent_within_the_code_of_another_command_do_not_publis
             result.Events.Should().NotContain(e => e is ReturnValueProduced);
         }
 
+        [Fact]
+        public async Task Commands_sent_via_API_within_a_split_submission_do_not_bubble_events()
+        {
+            using var kernel = new CompositeKernel
+            {
+                new CSharpKernel("cs1"),
+                new CSharpKernel("cs2")
+            };
+
+            var command = new SubmitCode($"""
+                #!cs1
+                using {typeof(Kernel).Namespace};
+                using {typeof(KernelCommand).Namespace};
+                var result = await Kernel.Root.SendAsync(new SubmitCode("123.Display();\n456", "cs2"));
+
+                #!cs2
+                Console.WriteLine(789);
+                """);
+
+            var result = await kernel.SendAsync(command);
+
+            using var _ = new AssertionScope();
+            result.Events.Should().NotContainErrors();
+            result.Events.Should().NotContain(e => e is DisplayedValueProduced);
+            result.Events.Should().NotContain(e => e is ReturnValueProduced);
+        }
+
         [Fact]
         public async Task Commands_sent_within_the_code_of_another_command_publish_CommandSucceeded_to_the_inner_result()
         {

src/Microsoft.DotNet.Interactive/Commands/KernelCommand.cs

@@ -14,6 +14,7 @@ public abstract class KernelCommand : IEquatable<KernelCommand>
 {
     private KernelCommand _parent;
     private string _token;
+    private List<KernelCommand> _childCommandsToBubbleEventsFrom;
 
     protected KernelCommand(
         string targetKernelName = null)
@@ -23,7 +24,8 @@ protected KernelCommand(
         RoutingSlip = new CommandRoutingSlip();
     }
 
-    [JsonIgnore] public KernelCommandInvocation Handler { get; set; }
+    [JsonIgnore] 
+    public KernelCommandInvocation Handler { get; set; }
 
     [JsonIgnore]
     public KernelCommand Parent => _parent;
@@ -115,17 +117,55 @@ static string CreateRootToken()
     private static int _nextRootToken = 0;
 #endif
 
-    [JsonIgnore] internal SchedulingScope SchedulingScope { get; set; }
+    [JsonIgnore] 
+    internal SchedulingScope SchedulingScope { get; set; }
 
-    [JsonIgnore] internal bool? ShouldPublishCompletionEvent { get; set; }
+    [JsonIgnore] 
+    internal bool? ShouldPublishCompletionEvent { get; set; }
 
     [JsonIgnore]
     public ParseResult KernelChooserParseResult { get; internal set; }
 
     [JsonIgnore]
     public CommandRoutingSlip RoutingSlip { get; }
 
-    internal bool ShouldResultIncludeEventsFromChildren { get; set; }
+    internal bool WasProxied { get; set; }
+
+    internal void ResultShouldIncludeEventsFrom(KernelCommand childCommand)
+    {
+        if (_childCommandsToBubbleEventsFrom is null)
+        {
+            _childCommandsToBubbleEventsFrom = new();
+        }
+
+        _childCommandsToBubbleEventsFrom.Add(childCommand);
+    }
+
+    internal bool ShouldResultIncludeEventsFrom(KernelCommand childCommand)
+    {
+        if (_childCommandsToBubbleEventsFrom is null)
+        {
+            return false;
+        }
+
+        for (var i = 0; i < _childCommandsToBubbleEventsFrom.Count; i++)
+        {
+            var command = _childCommandsToBubbleEventsFrom[i];
+
+            if (command.Equals(childCommand))
+            {
+                return true;
+            }
+
+            if (command.WasProxied &&
+                command.IsSelfOrDescendantOf(this))
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
 
     public virtual Task InvokeAsync(KernelInvocationContext context)
     {

src/Microsoft.DotNet.Interactive/Commands/RequestKernelInfo.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
-using System.Diagnostics;
 using System.Text.Json.Serialization;
 
 namespace Microsoft.DotNet.Interactive.Commands;

src/Microsoft.DotNet.Interactive/CompositeKernel.cs

@@ -190,13 +190,12 @@ private protected override async Task HandleRequestKernelInfoAsync(
     {
         context.Publish(new KernelInfoProduced(KernelInfo, command));
 
-        command.ShouldResultIncludeEventsFromChildren = true;
-
         foreach (var childKernel in ChildKernels)
         {
             if (childKernel.SupportsCommand(command))
             {
                 var childCommand = new RequestKernelInfo(childKernel.Name);
+                command.ResultShouldIncludeEventsFrom(childCommand);
                 childCommand.SetParent(command);
                 childCommand.RoutingSlip.ContinueWith(command.RoutingSlip);
                 await childKernel.HandleAsync(childCommand, context);

src/Microsoft.DotNet.Interactive/Connection/ProxyKernel.cs

@@ -73,6 +73,8 @@ private void UpdateKernelInfoFromEvent(KernelInfoProduced kernelInfoProduced)
 
     private Task HandleByForwardingToRemoteAsync(KernelCommand command, KernelInvocationContext context)
     {
+        command.WasProxied = true;
+
         if (command.OriginUri is null)
         {
             if (context.HandlingKernel == this)
@@ -98,6 +100,7 @@ private Task HandleByForwardingToRemoteAsync(KernelCommand command, KernelInvoca
                 return Task.CompletedTask;
             }
         }
+
         var targetKernelName = command.TargetKernelName;
         if (command.TargetKernelName == Name)
         {
@@ -127,7 +130,6 @@ private Task HandleByForwardingToRemoteAsync(KernelCommand command, KernelInvoca
         {
             command.TargetKernelName = targetKernelName;
 
-
             if (te.Result is CommandFailed cf)
             {
                 context.Fail(command, cf.Exception, cf.Message);

src/Microsoft.DotNet.Interactive/Kernel.Static.cs

@@ -112,9 +112,9 @@ public static void Javascript(string scriptContent)
 
         Task.Run(async () =>
         {
-            var kernelCommand = new DisplayValue(formatted);
-            context.Command.ShouldResultIncludeEventsFromChildren = true;
-            await kernel.SendAsync(kernelCommand);
+            var displayValue = new DisplayValue(formatted);
+            context.Command.ResultShouldIncludeEventsFrom(displayValue);
+            await kernel.SendAsync(displayValue);
         }).Wait(context.CancellationToken);
     }
 

src/Microsoft.DotNet.Interactive/Kernel.cs

@@ -161,8 +161,6 @@ private bool TrySplitCommand(
                     return false;
                 }
 
-                originalCommand.ShouldResultIncludeEventsFromChildren = true;
-
                 commands = new[] { adjustedCommand };
                 break;
 
@@ -171,13 +169,11 @@ private bool TrySplitCommand(
                 break;
         }
 
-        // FIX: (TrySplitCommand) 
-
         foreach (var command in commands)
         {
             if (!command.Equals(originalCommand))
             {
-                 originalCommand.ShouldResultIncludeEventsFromChildren = true;
+                 originalCommand.ResultShouldIncludeEventsFrom(command);
             }
 
             var handlingKernel = GetHandlingKernel(command, context);

src/Microsoft.DotNet.Interactive/KernelCommandResult.cs

@@ -24,10 +24,9 @@ public KernelCommandResult(KernelCommand command)
 
     internal void AddEvent(KernelEvent @event)
     {
-        // FIX: (AddEvent) 
         if (!@event.Command.Equals(Command))
         {
-            if (!Command.ShouldResultIncludeEventsFromChildren)
+            if (!Command.ShouldResultIncludeEventsFrom(@event.Command))
             {
                 return;
             }

src/Microsoft.DotNet.Interactive/KernelInvocationContext.cs

@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.ComponentModel;
 using System.Reactive.Linq;
 using System.Reactive.Subjects;
 using System.Threading;