From c864ff4ad6679b2c2d35f5909ac70f7d9b34dd9a Mon Sep 17 00:00:00 2001
From: reduckted <reduckted@outlook.com>
Date: Thu, 26 Jan 2023 12:00:38 +1000
Subject: [PATCH] Allow command interceptors to be unregistered.

---
 .../Commands/Commands.cs                      | 17 ++++++++-----
 .../Helpers/Disposable.cs                     | 24 +++++++++++++++++++
 .../VSSDK.Helpers.Shared.projitems            |  1 +
 3 files changed, 36 insertions(+), 6 deletions(-)
 create mode 100644 src/toolkit/Community.VisualStudio.Toolkit.Shared/Helpers/Disposable.cs

diff --git a/src/toolkit/Community.VisualStudio.Toolkit.Shared/Commands/Commands.cs b/src/toolkit/Community.VisualStudio.Toolkit.Shared/Commands/Commands.cs
index e8b7ad2..a7e3072 100644
--- a/src/toolkit/Community.VisualStudio.Toolkit.Shared/Commands/Commands.cs
+++ b/src/toolkit/Community.VisualStudio.Toolkit.Shared/Commands/Commands.cs
@@ -5,7 +5,6 @@
 using Microsoft.VisualStudio.OLE.Interop;
 using Microsoft.VisualStudio.Shell;
 using Microsoft.VisualStudio.Shell.Interop;
-using Task = System.Threading.Tasks.Task;
 
 namespace Community.VisualStudio.Toolkit
 {
@@ -93,32 +92,38 @@ public Task<bool> ExecuteAsync(CommandID cmd, string argument = "")
         /// <summary>
         /// Intercept any command before it is being handled by other command handlers.
         /// </summary>
-        public Task InterceptAsync(VSConstants.VSStd97CmdID command, Func<CommandProgression> func)
+        /// <returns>Returns an <see cref="IDisposable"/> that will remove the command interceptor when disposed.</returns>
+        public Task<IDisposable> InterceptAsync(VSConstants.VSStd97CmdID command, Func<CommandProgression> func)
             => InterceptAsync(typeof(VSConstants.VSStd97CmdID).GUID, (int)command, func);
 
         /// <summary>
         /// Intercept any command before it is being handled by other command handlers.
         /// </summary>
-        public Task InterceptAsync(VSConstants.VSStd2KCmdID command, Func<CommandProgression> func)
+        /// <returns>Returns an <see cref="IDisposable"/> that will remove the command interceptor when disposed.</returns>
+        public Task<IDisposable> InterceptAsync(VSConstants.VSStd2KCmdID command, Func<CommandProgression> func)
             => InterceptAsync(typeof(VSConstants.VSStd2KCmdID).GUID, (int)command, func);
 
         /// <summary>
         /// Intercept any command before it is being handled by other command handlers.
         /// </summary>
-        public Task InterceptAsync(Guid menuGroup, int commandId, Func<CommandProgression> func)
+        /// <returns>Returns an <see cref="IDisposable"/> that will remove the command interceptor when disposed.</returns>
+        public Task<IDisposable> InterceptAsync(Guid menuGroup, int commandId, Func<CommandProgression> func)
             => InterceptAsync(new CommandID(menuGroup, commandId), func);
 
         /// <summary>
         /// Intercept any command before it is being handled by other command handlers.
         /// </summary>
-        public async Task InterceptAsync(CommandID cmd, Func<CommandProgression> func)
+        /// <returns>Returns an <see cref="IDisposable"/> that will remove the command interceptor when disposed.</returns>
+        public async Task<IDisposable> InterceptAsync(CommandID cmd, Func<CommandProgression> func)
         {
             await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync();
 
             IVsRegisterPriorityCommandTarget? priority = await VS.Services.GetPriorityCommandTargetAsync();
             CommandInterceptor interceptor = new(cmd, func);
 
-            ErrorHandler.ThrowOnFailure(priority.RegisterPriorityCommandTarget(0, interceptor, out _));
+            ErrorHandler.ThrowOnFailure(priority.RegisterPriorityCommandTarget(0, interceptor, out uint cookie));
+
+            return new Disposable(() => priority.UnregisterPriorityCommandTarget(cookie));
         }
     }
 
diff --git a/src/toolkit/Community.VisualStudio.Toolkit.Shared/Helpers/Disposable.cs b/src/toolkit/Community.VisualStudio.Toolkit.Shared/Helpers/Disposable.cs
new file mode 100644
index 0000000..f484045
--- /dev/null
+++ b/src/toolkit/Community.VisualStudio.Toolkit.Shared/Helpers/Disposable.cs
@@ -0,0 +1,24 @@
+using System;
+
+namespace Community.VisualStudio.Toolkit
+{
+    internal sealed class Disposable : IDisposable
+    {
+        private readonly Action _action;
+        private bool _disposed;
+
+        public Disposable(Action action)
+        {
+            _action = action;
+        }
+
+        public void Dispose()
+        {
+            if (!_disposed)
+            {
+                _disposed = true;
+                _action();
+            }
+        }
+    }
+}
diff --git a/src/toolkit/Community.VisualStudio.Toolkit.Shared/VSSDK.Helpers.Shared.projitems b/src/toolkit/Community.VisualStudio.Toolkit.Shared/VSSDK.Helpers.Shared.projitems
index 968b78f..e5c6c98 100644
--- a/src/toolkit/Community.VisualStudio.Toolkit.Shared/VSSDK.Helpers.Shared.projitems
+++ b/src/toolkit/Community.VisualStudio.Toolkit.Shared/VSSDK.Helpers.Shared.projitems
@@ -10,6 +10,7 @@
   </PropertyGroup>
   <ItemGroup>
     <Compile Include="$(MSBuildThisFileDirectory)Debugger\Debugger.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Helpers\Disposable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Windows\IToolWindowPaneAware.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Windows\ToolkitToolWindowPane.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)\Attributes\ProvideBraceCompletionAttribute.cs" />