From 5e65a65561fd11924241b63bfc02662521b336ae Mon Sep 17 00:00:00 2001
From: Foahh <160589889+Foahh@users.noreply.github.com>
Date: Fri, 24 Oct 2025 17:44:43 +0800
Subject: [PATCH] fix(dialog): invoke the callbacks before cleanup process

---
 src/ShadUI/Controls/Dialog/DialogManager.cs | 28 ++++++++++-----------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/src/ShadUI/Controls/Dialog/DialogManager.cs b/src/ShadUI/Controls/Dialog/DialogManager.cs
index bb8f51e3..e3ddfd5a 100644
--- a/src/ShadUI/Controls/Dialog/DialogManager.cs
+++ b/src/ShadUI/Controls/Dialog/DialogManager.cs
@@ -52,18 +52,6 @@ internal void Show(Control control, DialogOptions options)
 
     internal void CloseDialog(Control control)
     {
-        var context = control.DataContext;
-        if (context is not null)
-        {
-            var contextType = context.GetType();
-            OnSuccessCallbacks.Remove(contextType);
-            OnSuccessWithContextCallbacks.Remove(contextType);
-            OnSuccessAsyncCallbacks.Remove(contextType);
-            OnSuccessWithContextAsyncCallbacks.Remove(contextType);
-            OnCancelCallbacks.Remove(contextType);
-            OnCancelAsyncCallbacks.Remove(contextType);
-        }
-
         Dialogs.Remove(control);
 
         OnDialogClosed?.Invoke(this, new DialogClosedEventArgs
@@ -96,7 +84,7 @@ internal void RemoveLast()
 
         var context = lastDialog.Key.DataContext;
         var contextType = context?.GetType();
-        if (contextType is not null) InvokeCallBacks(context, contextType, false);
+        if (contextType is not null) InvokeCallBacks(context, contextType, false, true);
     }
 
     internal readonly Dictionary<Type, Type> CustomDialogs = [];
@@ -120,7 +108,7 @@ public DialogManager Register<TView, TContext>() where TView : Control
     internal readonly Dictionary<Type, Func<Task>> OnSuccessAsyncCallbacks = [];
     internal readonly Dictionary<Type, Func<object, Task>> OnSuccessWithContextAsyncCallbacks = [];
 
-    private void InvokeCallBacks(object? context, Type type, bool success)
+    private void InvokeCallBacks(object? context, Type type, bool success, bool cleanup)
     {
         if (OnSuccessCallbacks.Remove(type, out var successCallback) && success)
         {
@@ -151,6 +139,16 @@ private void InvokeCallBacks(object? context, Type type, bool success)
         {
             cancelAsyncCallback?.Invoke();
         }
+
+        if (cleanup)
+        {
+            OnSuccessCallbacks.Remove(type);
+            OnSuccessWithContextCallbacks.Remove(type);
+            OnSuccessAsyncCallbacks.Remove(type);
+            OnSuccessWithContextAsyncCallbacks.Remove(type);
+            OnCancelCallbacks.Remove(type);
+            OnCancelAsyncCallbacks.Remove(type);
+        }
     }
 
     /// <summary>
@@ -169,7 +167,7 @@ public void Close<TContext>(TContext context, CloseDialogOptions? options = null
         foreach (var dialog in dialogs) CloseDialog(dialog.Key);
 
         var success = options?.Success ?? false;
-        InvokeCallBacks(context, typeof(TContext), success);
+        InvokeCallBacks(context, typeof(TContext), success, true);
 
         if (!clearAll) OpenLast();
     }