fix: clean up thread correctly to avoid memory leak

solara/server/kernel_context.py

@@ -311,6 +311,11 @@ def set_context_for_thread(context: VirtualKernelContext, thread: threading.Thre
     current_context[key] = context
 
 
+def clear_context_for_thread(thread: threading.Thread):
+    key = get_thread_key(thread)
+    current_context.pop(key, None)
+
+
 def has_current_context() -> bool:
     thread_key = get_current_thread_key()
     return (thread_key in current_context) and (current_context[thread_key] is not None)

solara/server/patch.py

@@ -259,18 +259,23 @@ def WidgetContextAwareThread__init__(self, *args, **kwargs):
 
 def Thread_debug_run(self):
     if not hasattr(self, "current_context"):
+        # this happens when a thread was running before we patched
         return Thread__run(self)
     if self.current_context:
         kernel_context.set_context_for_thread(self.current_context, self)
-        shell = self.current_context.kernel.shell
-        shell.display_pub.register_hook(shell.display_in_reacton_hook)
+        display_pub = self.current_context.kernel.shell.display_pub
+        display_in_reacton_hook = self.current_context.kernel.shell.display_in_reacton_hook
+        display_pub.register_hook(display_in_reacton_hook)
     try:
         context = self.current_context or solara.util.nullcontext()
         with pdb_guard(), context:
             Thread__run(self)
     finally:
-        if self.current_context:
-            shell.display_pub.unregister_hook(shell.display_in_reacton_hook)
+        current_context = self.current_context
+        self.current_context = None
+        kernel_context.clear_context_for_thread(self)
+        if current_context:
+            display_pub.unregister_hook(display_in_reacton_hook)
 
 
 _patched = False