diff --git a/Editor/Analysis/Action.cs b/Editor/Analysis/Action.cs
index 3f945a6b..4485cba3 100644
--- a/Editor/Analysis/Action.cs
+++ b/Editor/Analysis/Action.cs
@@ -385,7 +385,7 @@ public SyntaxNode GetRegistrationSyntax(string dialogueRunnerVariableName = "dia
                 return SyntaxFactory.ParseTypeName(t.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat));
             });
 
-            // If this is a function, or a command whose method returns a value, we also need to include the return type in
+            // If this is a function, we also need to include the return type in
             // this list.
             if (Type == ActionType.Function)
             {
@@ -393,14 +393,6 @@ public SyntaxNode GetRegistrationSyntax(string dialogueRunnerVariableName = "dia
 
                 typeArguments = typeArguments.Append(SyntaxFactory.ParseTypeName(returnType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)));
             }
-            if (Type == ActionType.Command) {
-                var returnType = (MethodSymbol as IMethodSymbol)?.ReturnType;
-
-                // If the command has a return type, and it's not void, add it to the list of type parameters
-                if (returnType != null && returnType.SpecialType != SpecialType.System_Void) {
-                    typeArguments = typeArguments.Append(SyntaxFactory.ParseTypeName(returnType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)));
-                }
-            }
 
             if (typeArguments.Any() && MethodSymbol?.IsStatic == true)
             {
diff --git a/Runtime/Commands/DefaultActions.cs b/Runtime/Commands/DefaultActions.cs
index 7dacef33..041fefac 100644
--- a/Runtime/Commands/DefaultActions.cs
+++ b/Runtime/Commands/DefaultActions.cs
@@ -31,7 +31,7 @@ public static void AddRegisterFunction() {
         public static void RegisterActions(IActionRegistration target, RegistrationType registrationType)
         {
             // Register the built-in methods and commands from Yarn Spinner for Unity.
-            target.AddCommandHandler<float, IEnumerator>("wait", Wait);
+            target.AddCommandHandler<float>("wait", Wait);
         }
 
         #region Commands
diff --git a/Runtime/IActionRegistration.cs b/Runtime/IActionRegistration.cs
index 6e74c00c..e64f1900 100644
--- a/Runtime/IActionRegistration.cs
+++ b/Runtime/IActionRegistration.cs
@@ -88,30 +88,7 @@ public interface IActionRegistration {
     }
 
     public static class ActionRegistrationExtension {
-        // GYB10 START
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<TResult>(this IActionRegistration registration, string commandName, System.Func<TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, T4, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, T4, T5, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, T9, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
-        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, TResult>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, TResult> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
-        // GYB10 END
+        
 
         /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
         public static void AddCommandHandler(this IActionRegistration registration, string commandName, System.Action handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
@@ -139,6 +116,54 @@ public static class ActionRegistrationExtension {
         public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10>(this IActionRegistration registration, string commandName, System.Action<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
         // GYB11 END
        
+        public static void AddCommandHandler<T1>(this IActionRegistration registration, string commandName, System.Func<T1, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2>(this IActionRegistration registration, string commandName, System.Func<T1, T2, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, T9>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, System.Threading.Tasks.Task> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+       
+       #if USE_UNITASK
+       
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<(this IActionRegistration registration, string commandName, System.Func<Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1>(this IActionRegistration registration, string commandName, System.Func<T1, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2>(this IActionRegistration registration, string commandName, System.Func<T1, T2, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, T9>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
+        public static void AddCommandHandler<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10>(this IActionRegistration registration, string commandName, System.Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, Cysharp.Threading.Tasks.UniTask> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
+        // GYB11 END
+
+        #endif
+       
         public static void AddCommandHandler<T1>(this IActionRegistration registration, string commandName, System.Func<T1, IEnumerator> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
         /// <inheritdoc cref="AddCommandHandler(string, Delegate)"/>
         public static void AddCommandHandler<T1, T2>(this IActionRegistration registration, string commandName, System.Func<T1, T2, IEnumerator> handler) => registration.AddCommandHandler(commandName, (Delegate)handler);
diff --git a/SourceGenerator/YarnSpinner.Unity.SourceCodeGenerator.dll b/SourceGenerator/YarnSpinner.Unity.SourceCodeGenerator.dll
index 3d75609d..0de2c25f 100644
Binary files a/SourceGenerator/YarnSpinner.Unity.SourceCodeGenerator.dll and b/SourceGenerator/YarnSpinner.Unity.SourceCodeGenerator.dll differ
diff --git a/Tests/Runtime/DialogueViewTests/LineViewTests.cs b/Tests/Runtime/DialogueViewTests/LineViewTests.cs
index 053c8c85..12edf87e 100644
--- a/Tests/Runtime/DialogueViewTests/LineViewTests.cs
+++ b/Tests/Runtime/DialogueViewTests/LineViewTests.cs
@@ -11,11 +11,11 @@ namespace Yarn.Unity.Tests
     using UnityEngine;
     using UnityEngine.TestTools;
     using UnityEngine.SceneManagement;
+    using System;
 
 #if USE_UNITASK
     using Cysharp.Threading.Tasks;
     using YarnTask = Cysharp.Threading.Tasks.UniTask;
-    using System;
 #else
     using YarnTask = System.Threading.Tasks.Task;
 #endif
diff --git a/Tests/Runtime/EndToEndTests/EndToEndUtility.cs b/Tests/Runtime/EndToEndTests/EndToEndUtility.cs
index 7d6db163..8b5c4e1f 100644
--- a/Tests/Runtime/EndToEndTests/EndToEndUtility.cs
+++ b/Tests/Runtime/EndToEndTests/EndToEndUtility.cs
@@ -145,7 +145,7 @@ await WaitForTaskAsync(
         public static async YarnTask WaitForTaskAsync(YarnTask task, string? failureMessage = null, int timeoutMilliseconds = 2000)
         {
             try {
-                await task.Wait(TimeSpan.FromMilliseconds(timeoutMilliseconds));
+                await YarnAsync.Wait(task, TimeSpan.FromMilliseconds(timeoutMilliseconds));
             } catch (TimeoutException timeout) {
                 if (failureMessage == null) {
                     throw;