Making the raw TokenSource classes available to coroutines and unitasks

Runtime/Scripts/TokenSource/TokenSource.cs

@@ -3,6 +3,7 @@
 using System.Linq;
 using System.Net.Http;
 using System.Threading.Tasks;
+using LiveKit.Internal.Threading;
 using Newtonsoft.Json;
 
 namespace LiveKit
@@ -21,7 +22,7 @@ public interface ITokenSource
     /// </summary>
     public interface ITokenSourceFixed : ITokenSource
     {
-        public Task<ConnectionDetails> FetchConnectionDetails();
+        public TaskYieldInstruction<ConnectionDetails> FetchConnectionDetails();
     }
 
     /// <summary>
@@ -30,7 +31,7 @@ public interface ITokenSourceFixed : ITokenSource
     /// </summary>
     public interface ITokenSourceConfigurable : ITokenSource
     {
-        public Task<ConnectionDetails> FetchConnectionDetails(TokenSourceFetchOptions options);
+        public TaskYieldInstruction<ConnectionDetails> FetchConnectionDetails(TokenSourceFetchOptions options);
     }
 
     /// <summary>
@@ -48,10 +49,10 @@ public TokenSourceLiteral(string serverUrl, string participantToken)
             _participantToken = participantToken;
         }
 
-        public Task<ConnectionDetails> FetchConnectionDetails()
+        public TaskYieldInstruction<ConnectionDetails> FetchConnectionDetails()
         {
             var result = new ConnectionDetails { ServerUrl = _serverUrl, ParticipantToken = _participantToken };
-            return Task.FromResult(result);
+            return new TaskYieldInstruction<ConnectionDetails>(Task.FromResult(result));
         }
     }
 
@@ -70,9 +71,22 @@ public TokenSourceCustom(CustomTokenFunction customTokenFunction)
             _customTokenFunction = customTokenFunction;
         }
 
-        public Task<ConnectionDetails> FetchConnectionDetails()
+        public TaskYieldInstruction<ConnectionDetails> FetchConnectionDetails()
         {
-            return _customTokenFunction();
+            // Route a synchronous throw (or a null return) from the user's function through the
+            // instruction's IsError/Exception, so callers never have to guard the call itself.
+            Task<ConnectionDetails> task;
+            try
+            {
+                task = _customTokenFunction()
+                    ?? Task.FromException<ConnectionDetails>(
+                        new InvalidOperationException("Custom token function returned a null task"));
+            }
+            catch (Exception e)
+            {
+                task = Task.FromException<ConnectionDetails>(e);
+            }
+            return new TaskYieldInstruction<ConnectionDetails>(task);
         }
     }
 
@@ -94,7 +108,14 @@ public TokenSourceEndpoint(string endpointUrl, IEnumerable<StringPair> headers)
             _headers = headers;
         }
 
-        public async Task<ConnectionDetails> FetchConnectionDetails(TokenSourceFetchOptions options)
+        public TaskYieldInstruction<ConnectionDetails> FetchConnectionDetails(TokenSourceFetchOptions options)
+        {
+            // Async methods can't return the (non-awaitable) instruction directly, so the actual
+            // request lives in the helper below; the returned task carries any synchronous throw.
+            return new TaskYieldInstruction<ConnectionDetails>(FetchConnectionDetailsAsync(options));
+        }
+
+        private async Task<ConnectionDetails> FetchConnectionDetailsAsync(TokenSourceFetchOptions options)
         {
             var requestBody = BuildRequest(options);
             var jsonBody = JsonConvert.SerializeObject(requestBody);

Runtime/Scripts/TokenSource/TokenSourceComponent.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Threading.Tasks;
 using UnityEngine;
 using LiveKit.Internal;
 
@@ -60,23 +59,19 @@ public void Awake()
         /// </summary>
         public TaskYieldInstruction<ConnectionDetails> FetchConnectionDetails(TokenSourceFetchOptions options)
         {
-            Task<ConnectionDetails> task;
             switch (_tokenSource)
             {
                 case ITokenSourceConfigurable configurableSource:
-                    task = configurableSource.FetchConnectionDetails(Coalesce(_config, options));
-                    break;
+                    return configurableSource.FetchConnectionDetails(Coalesce(_config, options));
 
                 case ITokenSourceFixed fixedSource:
                     if (options != null)
                         Utils.Warning("TokenSourceComponent uses a fixed config, so fetch options are ignored.");
-                    task = fixedSource.FetchConnectionDetails();
-                    break;
+                    return fixedSource.FetchConnectionDetails();
 
                 default:
                     throw new InvalidOperationException("Unknown token source type");
             }
-            return new TaskYieldInstruction<ConnectionDetails>(task);
         }
 
         private static TokenSourceFetchOptions Coalesce(TokenSourceComponentConfig config, TokenSourceFetchOptions? options)