From d47f56c4052c202da51ab631485da8a72ca33af7 Mon Sep 17 00:00:00 2001
From: gchow <garson.chow@factset.com>
Date: Thu, 11 Jan 2024 12:44:36 -0500
Subject: [PATCH] Reflect theme changes without a refresh

By leveraging `IThemeManager.themeChanged`, we can listen on theme change signals
and rebuild the theme object in response. This allows CSS variable changes to reflect
in the MUI theme without having to refresh the page.
---
 .../src/components/jl-theme-provider.tsx           |  7 ++++++-
 packages/jupyter-ai/src/index.ts                   | 14 +++++++++++++-
 packages/jupyter-ai/src/theme-provider.ts          | 11 +++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/packages/jupyter-ai/src/components/jl-theme-provider.tsx b/packages/jupyter-ai/src/components/jl-theme-provider.tsx
index 7e6fc2113..239c410f7 100644
--- a/packages/jupyter-ai/src/components/jl-theme-provider.tsx
+++ b/packages/jupyter-ai/src/components/jl-theme-provider.tsx
@@ -1,7 +1,7 @@
 import React, { useState, useEffect } from 'react';
 
 import { Theme, ThemeProvider, createTheme } from '@mui/material/styles';
-import { getJupyterLabTheme } from '../theme-provider';
+import { getJupyterLabTheme, getThemeManager } from '../theme-provider';
 
 export function JlThemeProvider(props: {
   children: React.ReactNode;
@@ -13,6 +13,11 @@ export function JlThemeProvider(props: {
       setTheme(await getJupyterLabTheme());
     }
     setJlTheme();
+
+    const manager = getThemeManager();
+    if (manager) {
+      manager.themeChanged.connect(setJlTheme);
+    }
   }, []);
 
   return <ThemeProvider theme={theme}>{props.children}</ThemeProvider>;
diff --git a/packages/jupyter-ai/src/index.ts b/packages/jupyter-ai/src/index.ts
index e48e2b211..5447e6774 100644
--- a/packages/jupyter-ai/src/index.ts
+++ b/packages/jupyter-ai/src/index.ts
@@ -4,7 +4,11 @@ import {
   ILayoutRestorer
 } from '@jupyterlab/application';
 
-import { IWidgetTracker, ReactWidget } from '@jupyterlab/apputils';
+import {
+  IWidgetTracker,
+  ReactWidget,
+  IThemeManager
+} from '@jupyterlab/apputils';
 import { IDocumentWidget } from '@jupyterlab/docregistry';
 import { IGlobalAwareness } from '@jupyter/collaboration';
 import type { Awareness } from 'y-protocols/awareness';
@@ -12,6 +16,7 @@ import { buildChatSidebar } from './widgets/chat-sidebar';
 import { SelectionWatcher } from './selection-watcher';
 import { ChatHandler } from './chat_handler';
 import { buildErrorWidget } from './widgets/chat-error';
+import { setThemeManager } from './theme-provider';
 
 export type DocumentTracker = IWidgetTracker<IDocumentWidget>;
 
@@ -22,8 +27,10 @@ const plugin: JupyterFrontEndPlugin<void> = {
   id: 'jupyter_ai:plugin',
   autoStart: true,
   optional: [IGlobalAwareness, ILayoutRestorer],
+  requires: [IThemeManager],
   activate: async (
     app: JupyterFrontEnd,
+    manager: IThemeManager,
     globalAwareness: Awareness | null,
     restorer: ILayoutRestorer | null
   ) => {
@@ -57,6 +64,11 @@ const plugin: JupyterFrontEndPlugin<void> = {
     if (restorer) {
       restorer.add(chatWidget, 'jupyter-ai-chat');
     }
+
+    /**
+     * Set the theme manager
+     */
+    setThemeManager(manager);
   }
 };
 
diff --git a/packages/jupyter-ai/src/theme-provider.ts b/packages/jupyter-ai/src/theme-provider.ts
index 405f08198..cb503b3ad 100644
--- a/packages/jupyter-ai/src/theme-provider.ts
+++ b/packages/jupyter-ai/src/theme-provider.ts
@@ -1,5 +1,8 @@
+import { IThemeManager } from '@jupyterlab/apputils';
 import { Theme, createTheme } from '@mui/material/styles';
 
+let themeManager: IThemeManager | null = null;
+
 function getCSSVariable(name: string): string {
   return getComputedStyle(document.body).getPropertyValue(name).trim();
 }
@@ -10,6 +13,14 @@ export async function pollUntilReady(): Promise<void> {
   }
 }
 
+export function setThemeManager(manager: IThemeManager): void {
+  themeManager = manager;
+}
+
+export function getThemeManager(): IThemeManager | null {
+  return themeManager;
+}
+
 export async function getJupyterLabTheme(): Promise<Theme> {
   await pollUntilReady();
   const light = document.body.getAttribute('data-jp-theme-light');