From ee5020c41f0bf776af18c5202814750ff029f591 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Mon, 15 Jan 2024 17:37:50 -0800
Subject: [PATCH] actually set task

---
 packages/core/src/scheduler/task.rs | 2 +-
 packages/core/src/scheduler/wait.rs | 2 ++
 packages/core/src/scope_context.rs  | 7 +++----
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/packages/core/src/scheduler/task.rs b/packages/core/src/scheduler/task.rs
index 6cddf8043..7529b49cf 100644
--- a/packages/core/src/scheduler/task.rs
+++ b/packages/core/src/scheduler/task.rs
@@ -89,7 +89,7 @@ impl Runtime {
     /// Drop the future with the given TaskId
     ///
     /// This does not abort the task, so you'll want to wrap it in an abort handle if that's important to you
-    pub fn remove(&self, id: Task) -> Option<LocalTask> {
+    pub(crate) fn remove_task(&self, id: Task) -> Option<LocalTask> {
         self.tasks.borrow_mut().try_remove(id.0)
     }
 
diff --git a/packages/core/src/scheduler/wait.rs b/packages/core/src/scheduler/wait.rs
index 45b9c3791..2f208a525 100644
--- a/packages/core/src/scheduler/wait.rs
+++ b/packages/core/src/scheduler/wait.rs
@@ -21,6 +21,7 @@ impl VirtualDom {
         // update the scope stack
         self.runtime.scope_stack.borrow_mut().push(task.scope);
         self.runtime.rendering.set(false);
+        self.runtime.current_task.set(Some(id));
 
         // If the task completes...
         if task.task.borrow_mut().as_mut().poll(&mut cx).is_ready() {
@@ -35,5 +36,6 @@ impl VirtualDom {
         // Remove the scope from the stack
         self.runtime.scope_stack.borrow_mut().pop();
         self.runtime.rendering.set(true);
+        self.runtime.current_task.set(None);
     }
 }
diff --git a/packages/core/src/scope_context.rs b/packages/core/src/scope_context.rs
index 659a08323..6c6526868 100644
--- a/packages/core/src/scope_context.rs
+++ b/packages/core/src/scope_context.rs
@@ -8,7 +8,6 @@ use std::{
     any::Any,
     cell::{Cell, RefCell},
     future::Future,
-    rc::Rc,
     sync::Arc,
 };
 
@@ -250,7 +249,7 @@ impl ScopeContext {
     ///
     /// This drops the task immediately.
     pub fn remove_future(&self, id: Task) {
-        with_runtime(|rt| rt.remove(id)).expect("Runtime to exist");
+        with_runtime(|rt| rt.remove_task(id)).expect("Runtime to exist");
     }
 
     /// Mark this component as suspended and then return None
@@ -314,10 +313,10 @@ impl ScopeContext {
 
 impl Drop for ScopeContext {
     fn drop(&mut self) {
+        // Drop all spawned tasks
         with_runtime(|rt| {
-            // Drop all spawned tasks
             for id in self.spawned_tasks.borrow().iter() {
-                rt.remove(*id);
+                rt.remove_task(*id);
             }
         })
         .expect("Runtime to exist")