Merge branch 'main' into copilot/update-event-handlers-to-concurrent-dictionary

Author: stephentoub

dotnet/src/ActionDisposable.cs

@@ -0,0 +1,24 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+namespace GitHub.Copilot.SDK;
+
+/// <summary>
+/// A disposable that invokes an action when disposed.
+/// </summary>
+internal sealed class ActionDisposable : IDisposable
+{
+    private Action? _action;
+
+    public ActionDisposable(Action action)
+    {
+        _action = action;
+    }
+
+    public void Dispose()
+    {
+        var action = Interlocked.Exchange(ref _action, null);
+        action?.Invoke();
+    }
+}

dotnet/src/Client.cs

@@ -1561,22 +1561,3 @@ public class ToolResultAIContent(ToolResultObject toolResult) : AIContent
 {
     public ToolResultObject Result => toolResult;
 }
-
-/// <summary>
-/// A disposable that invokes an action when disposed.
-/// </summary>
-internal sealed class ActionDisposable : IDisposable
-{
-    private Action? _action;
-
-    public ActionDisposable(Action action)
-    {
-        _action = action;
-    }
-
-    public void Dispose()
-    {
-        var action = Interlocked.Exchange(ref _action, null);
-        action?.Invoke();
-    }
-}

dotnet/src/Session.cs

@@ -48,10 +48,8 @@ public partial class CopilotSession : IAsyncDisposable
     private readonly ConcurrentDictionary<SessionEventHandler, bool> _eventHandlers = new();
     private readonly Dictionary<string, AIFunction> _toolHandlers = new();
     private readonly JsonRpc _rpc;
-    private PermissionRequestHandler? _permissionHandler;
-    private readonly SemaphoreSlim _permissionHandlerLock = new(1, 1);
-    private UserInputHandler? _userInputHandler;
-    private readonly SemaphoreSlim _userInputHandlerLock = new(1, 1);
+    private volatile PermissionRequestHandler? _permissionHandler;
+    private volatile UserInputHandler? _userInputHandler;
     private SessionHooks? _hooks;
     private readonly SemaphoreSlim _hooksLock = new(1, 1);
     private SessionRpc? _sessionRpc;
@@ -247,7 +245,7 @@ void Handler(SessionEvent evt)
     public IDisposable On(SessionEventHandler handler)
     {
         _eventHandlers.TryAdd(handler, true);
-        return new OnDisposeCall(() => _eventHandlers.TryRemove(handler, out _));
+        return new ActionDisposable(() => _eventHandlers.TryRemove(handler, out _));
     }
 
     /// <summary>
@@ -259,7 +257,7 @@ public IDisposable On(SessionEventHandler handler)
     /// </remarks>
     internal void DispatchEvent(SessionEvent sessionEvent)
     {
-        foreach (var (handler, _) in _eventHandlers)
+        foreach (var handler in _eventHandlers.Keys)
         {
             // We allow handler exceptions to propagate so they are not lost
             handler(sessionEvent);
@@ -301,15 +299,7 @@ internal void RegisterTools(ICollection<AIFunction> tools)
     /// </remarks>
     internal void RegisterPermissionHandler(PermissionRequestHandler handler)
     {
-        _permissionHandlerLock.Wait();
-        try
-        {
-            _permissionHandler = handler;
-        }
-        finally
-        {
-            _permissionHandlerLock.Release();
-        }
+        _permissionHandler = handler;
     }
 
     /// <summary>
@@ -319,16 +309,7 @@ internal void RegisterPermissionHandler(PermissionRequestHandler handler)
     /// <returns>A task that resolves with the permission decision.</returns>
     internal async Task<PermissionRequestResult> HandlePermissionRequestAsync(JsonElement permissionRequestData)
     {
-        await _permissionHandlerLock.WaitAsync();
-        PermissionRequestHandler? handler;
-        try
-        {
-            handler = _permissionHandler;
-        }
-        finally
-        {
-            _permissionHandlerLock.Release();
-        }
+        var handler = _permissionHandler;
 
         if (handler == null)
         {
@@ -355,15 +336,7 @@ internal async Task<PermissionRequestResult> HandlePermissionRequestAsync(JsonEl
     /// <param name="handler">The handler to invoke when user input is requested.</param>
     internal void RegisterUserInputHandler(UserInputHandler handler)
     {
-        _userInputHandlerLock.Wait();
-        try
-        {
-            _userInputHandler = handler;
-        }
-        finally
-        {
-            _userInputHandlerLock.Release();
-        }
+        _userInputHandler = handler;
     }
 
     /// <summary>
@@ -373,16 +346,7 @@ internal void RegisterUserInputHandler(UserInputHandler handler)
     /// <returns>A task that resolves with the user's response.</returns>
     internal async Task<UserInputResponse> HandleUserInputRequestAsync(UserInputRequest request)
     {
-        await _userInputHandlerLock.WaitAsync();
-        UserInputHandler? handler;
-        try
-        {
-            handler = _userInputHandler;
-        }
-        finally
-        {
-            _userInputHandlerLock.Release();
-        }
+        var handler = _userInputHandler;
 
         if (handler == null)
         {
@@ -590,20 +554,7 @@ await InvokeRpcAsync<object>(
         _eventHandlers.Clear();
         _toolHandlers.Clear();
 
-        await _permissionHandlerLock.WaitAsync();
-        try
-        {
-            _permissionHandler = null;
-        }
-        finally
-        {
-            _permissionHandlerLock.Release();
-        }
-    }
-
-    private class OnDisposeCall(Action callback) : IDisposable
-    {
-        public void Dispose() => callback();
+        _permissionHandler = null;
     }
 
     internal record SendMessageRequest