chore: implement polling toolchain task events instead of pushing

Assets/Rivet/Editor/RivetToolchain.cs

@@ -1,52 +1,159 @@
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.Runtime.InteropServices;
+using Newtonsoft.Json.Linq;
 
 namespace Rivet.Editor
 {
-    public static class RivetToolchain
+    internal static class RivetToolchainFFI
     {
         const string RustLibrary = "__Internal";
 
-        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
-        public delegate void EventCallback(ulong taskId, IntPtr eventJson);
+        [StructLayout(LayoutKind.Sequential)]
+        public struct RunTaskResult
+        {
+            public ulong TaskId;
+            public byte ErrorCode;
+        }
 
         [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_run_task")]
-        private static extern ulong run_task(
+        public static extern RunTaskResult run_task(
             [MarshalAs(UnmanagedType.LPStr)] string name,
-            [MarshalAs(UnmanagedType.LPStr)] string inputJson,
-            EventCallback callback
+            [MarshalAs(UnmanagedType.LPStr)] string inputJson
         );
 
         [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_abort_task")]
-        private static extern bool abort_task(ulong taskId);
+        public static extern byte abort_task(ulong taskId);
 
         [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_shutdown")]
-        private static extern void shutdown();
+        public static extern void shutdown();
 
         [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_free_rust_string")]
         private static extern void free_rust_string(IntPtr str);
 
-        public static ulong RunTask(string name, string inputJson, EventCallback callback)
+        public static string PtrToString(IntPtr ptr)
         {
-            return run_task(name, inputJson, callback);
+            string result = Marshal.PtrToStringAnsi(ptr);
+            free_rust_string(ptr);
+            return result;
         }
 
-        public static bool AbortTask(ulong taskId)
+        [StructLayout(LayoutKind.Sequential)]
+        public struct TaskEvent
         {
-            return abort_task(taskId);
+            public ulong TaskId;
+            public IntPtr EventJson;
+        }
+
+        [StructLayout(LayoutKind.Sequential)]
+        public struct PollTaskEventsResult
+        {
+            public UIntPtr Count;
+            public byte ErrorCode;
+        }
+
+        [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_poll_task_events")]
+        public static extern PollTaskEventsResult poll_task_events(
+            [In, Out] TaskEvent[] events,
+            UIntPtr maxEvents
+        );
+
+        public const uint POLL_TASK_EVENT_BUFFER = 128;
+    }
+
+    public static class RivetToolchain
+    {
+        private static Dictionary<ulong, Action<JObject>> taskCallbacks = new Dictionary<ulong, Action<JObject>>();
+        private static readonly object taskCallbacksLock = new object();
+
+        public static ulong RunTask(string name, string inputJson, Action<JObject> callback)
+        {
+            RivetToolchainFFI.RunTaskResult result = RivetToolchainFFI.run_task(name, inputJson);
+            if (result.ErrorCode != 0)
+            {
+                throw new Exception($"Rivet task '{name}' failed with error code: {result.ErrorCode}");
+            }
+
+            lock (taskCallbacksLock)
+            {
+                taskCallbacks[result.TaskId] = callback;
+            }
+
+            return result.TaskId;
+        }
+
+        public static void AbortTask(ulong taskId)
+        {
+            var error_code = RivetToolchainFFI.abort_task(taskId);
+            if (error_code != 0)
+            {
+                throw new Exception($"Failed to abort Rivet task with ID: {taskId}");
+            }
         }
 
         public static void Shutdown()
         {
-            // shutdown();
+            // RivetToolchainFFI.shutdown();
         }
 
-        public static string PtrToString(IntPtr ptr)
+        public class TaskEventData
         {
-            string result = Marshal.PtrToStringAnsi(ptr);
-            free_rust_string(ptr);
-            return result;
+            public ulong TaskId;
+            public JObject EventJson;
+        }
+
+        public static List<TaskEventData> PollTaskEvents()
+        {
+            List<TaskEventData> allEvents = new List<TaskEventData>();
+            RivetToolchainFFI.TaskEvent[] buffer = new RivetToolchainFFI.TaskEvent[RivetToolchainFFI.POLL_TASK_EVENT_BUFFER];
+
+            uint totalPolled = 0;
+            while (true)
+            {
+                RivetToolchainFFI.PollTaskEventsResult result = RivetToolchainFFI.poll_task_events(buffer, new UIntPtr(RivetToolchainFFI.POLL_TASK_EVENT_BUFFER));
+
+                if (result.ErrorCode != 0)
+                {
+                    throw new Exception($"Failed to poll Rivet task events with error code: {result.ErrorCode}");
+                }
+
+                uint polledCount = result.Count.ToUInt32();
+                totalPolled += polledCount;
+
+                for (int i = 0; i < polledCount; i++)
+                {
+                    string eventJsonString = RivetToolchainFFI.PtrToString(buffer[i].EventJson);
+                    JObject eventJson = JObject.Parse(eventJsonString);
+
+                    // Attempt to get callback
+                    Action<JObject> callback = null;
+                    lock (taskCallbacksLock)
+                    {
+                        if (!taskCallbacks.TryGetValue(buffer[i].TaskId, out callback))
+                        {
+                            RivetLogger.Warning($"Missing callback for task ID: {buffer[i].TaskId}");
+                        }
+                    }
+
+                    // Call callback
+                    callback?.Invoke(eventJson);
+                }
+
+                if (polledCount == 0)
+                {
+                    // No more events to poll
+                    break;
+                }
+            }
+
+            // Log task event
+            if (totalPolled > 0)
+            {
+                RivetLogger.Log($"Polled {totalPolled} task events");
+            }
+
+            return allEvents;
         }
     }
 }

Assets/Rivet/Editor/Task.cs

@@ -49,15 +49,8 @@ private void Run(string name, string inputJson)
             _taskId = RivetToolchain.RunTask(name, inputJson, OnOutputEvent);
         }
 
-        private void OnOutputEvent(ulong taskId, IntPtr eventJsonPtr)
+        private void OnOutputEvent(JObject eventObj)
         {
-            string eventJson = RivetToolchain.PtrToString(eventJsonPtr);
-            EditorApplication.delayCall += () => HandleOnOutputEvent(eventJson);
-        }
-
-        private void HandleOnOutputEvent(string eventJson)
-        {
-            var eventObj = JObject.Parse(eventJson);
             if (eventObj.ContainsKey("log"))
             {
                 OnLogEvent((string)eventObj["log"]);
@@ -85,7 +78,7 @@ private void HandleOnOutputEvent(string eventJson)
             }
             else
             {
-                RivetLogger.Warning($"Unknown event {eventJson}");
+                RivetLogger.Warning($"Unknown event {eventObj.ToString(Newtonsoft.Json.Formatting.None)}");
             }
         }
 

Assets/Rivet/Editor/UI/Dock/Dock.cs

@@ -181,6 +181,9 @@ public void OnEnable()
                 autoRestart: true
             );
 
+            // Add update callback
+            EditorApplication.update += OnEditorUpdate;
+
             // Shut down on reload
             AssemblyReloadEvents.beforeAssemblyReload += () =>
             {
@@ -199,6 +202,9 @@ public void OnEnable()
         public void OnDisable()
         {
             RivetLogger.Log("On Disable");
+
+            EditorApplication.update -= OnEditorUpdate;
+
             ShutdownPlugin();
         }
 
@@ -305,5 +311,10 @@ private void OpenHub()
                 Application.OpenURL($"https://hub.rivet.gg");
             }
         }
+
+        private void OnEditorUpdate()
+        {
+            RivetToolchain.PollTaskEvents();
+        }
     }
 }