migrate theme message protobus (RooCodeInc#4012)

Co-authored-by: Elephant Lumps <[email protected]>

Author: celestial-vault

proto/ui.proto

@@ -246,4 +246,7 @@ service UiService {
 
   // Subscribe to partial message updates (streaming Cline messages as they're built)
   rpc subscribeToPartialMessage(EmptyRequest) returns (stream ClineMessage);
+  
+  // Subscribe to theme change events
+  rpc subscribeToTheme(EmptyRequest) returns (stream String);
 }

src/core/controller/index.ts

@@ -224,12 +224,6 @@ export class Controller {
 			case "webviewDidLaunch":
 				this.postStateToWebview()
 				this.workspaceTracker?.populateFilePaths() // don't await
-				getTheme().then((theme) =>
-					this.postMessageToWebview({
-						type: "theme",
-						text: JSON.stringify(theme),
-					}),
-				)
 				// post last cached models in case the call to endpoint fails
 				this.readOpenRouterModels().then((openRouterModels) => {
 					if (openRouterModels) {

src/core/controller/ui/subscribeToTheme.ts

@@ -0,0 +1,76 @@
+import { Controller } from "../index"
+import { EmptyRequest, String } from "@shared/proto/common"
+import { StreamingResponseHandler, getRequestRegistry } from "../grpc-handler"
+import { getTheme } from "@integrations/theme/getTheme"
+
+// Keep track of active theme subscriptions
+const activeThemeSubscriptions = new Set<StreamingResponseHandler>()
+
+/**
+ * Subscribe to theme change events
+ * @param controller The controller instance
+ * @param request The empty request
+ * @param responseStream The streaming response handler
+ * @param requestId The ID of the request (passed by the gRPC handler)
+ */
+export async function subscribeToTheme(
+	controller: Controller,
+	request: EmptyRequest,
+	responseStream: StreamingResponseHandler,
+	requestId?: string,
+): Promise<void> {
+	// Add this subscription to the active subscriptions
+	activeThemeSubscriptions.add(responseStream)
+
+	// Register cleanup when the connection is closed
+	const cleanup = () => {
+		activeThemeSubscriptions.delete(responseStream)
+	}
+
+	// Register the cleanup function with the request registry if we have a requestId
+	if (requestId) {
+		getRequestRegistry().registerRequest(requestId, cleanup, { type: "theme_subscription" }, responseStream)
+	}
+
+	// Send the current theme immediately upon subscription
+	const theme = await getTheme()
+	if (theme) {
+		try {
+			const themeEvent = String.create({
+				value: JSON.stringify(theme),
+			})
+			await responseStream(
+				themeEvent,
+				false, // Not the last message
+			)
+		} catch (error) {
+			console.error("Error sending initial theme:", error)
+			activeThemeSubscriptions.delete(responseStream)
+		}
+	}
+}
+
+/**
+ * Send a theme event to all active subscribers
+ * @param themeJson The JSON-stringified theme data
+ */
+export async function sendThemeEvent(themeJson: string): Promise<void> {
+	// Send the event to all active subscribers
+	const promises = Array.from(activeThemeSubscriptions).map(async (responseStream) => {
+		try {
+			const event = String.create({
+				value: themeJson,
+			})
+			await responseStream(
+				event,
+				false, // Not the last message
+			)
+		} catch (error) {
+			console.error("Error sending theme event:", error)
+			// Remove the subscription if there was an error
+			activeThemeSubscriptions.delete(responseStream)
+		}
+	})
+
+	await Promise.all(promises)
+}

src/core/webview/index.ts

@@ -8,6 +8,7 @@ import { findLast } from "@shared/array"
 import { readFile } from "fs/promises"
 import path from "node:path"
 import { WebviewProviderType } from "@/shared/webview/types"
+import { sendThemeEvent } from "@core/controller/ui/subscribeToTheme"
 
 /*
 https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default/weather-webview/src/providers/WeatherViewProvider.ts
@@ -139,11 +140,11 @@ export class WebviewProvider implements vscode.WebviewViewProvider {
 			vscode.workspace.onDidChangeConfiguration(
 				async (e) => {
 					if (e && e.affectsConfiguration("workbench.colorTheme")) {
-						// Sends latest theme name to webview
-						await this.controller.postMessageToWebview({
-							type: "theme",
-							text: JSON.stringify(await getTheme()),
-						})
+						// Send theme update via gRPC subscription
+						const theme = await getTheme()
+						if (theme) {
+							await sendThemeEvent(JSON.stringify(theme))
+						}
 					}
 					if (e && e.affectsConfiguration("cline.mcpMarketplace.enabled")) {
 						// Update state when marketplace tab setting changes

src/shared/ExtensionMessage.ts

@@ -19,7 +19,6 @@ export interface ExtensionMessage {
 		| "selectedImages"
 		| "ollamaModels"
 		| "lmStudioModels"
-		| "theme"
 		| "workspaceUpdated"
 		| "openRouterModels"
 		| "openAiModels"

webview-ui/src/context/ExtensionStateContext.tsx

@@ -197,12 +197,6 @@ export const ExtensionStateContextProvider: React.FC<{
 	const handleMessage = useCallback((event: MessageEvent) => {
 		const message: ExtensionMessage = event.data
 		switch (message.type) {
-			case "theme": {
-				if (message.text) {
-					setTheme(convertTextMateToHljs(JSON.parse(message.text)))
-				}
-				break
-			}
 			case "workspaceUpdated": {
 				setFilePaths(message.filePaths ?? [])
 				break
@@ -246,6 +240,7 @@ export const ExtensionStateContextProvider: React.FC<{
 	const settingsButtonClickedSubscriptionRef = useRef<(() => void) | null>(null)
 	const partialMessageUnsubscribeRef = useRef<(() => void) | null>(null)
 	const mcpMarketplaceUnsubscribeRef = useRef<(() => void) | null>(null)
+	const themeSubscriptionRef = useRef<(() => void) | null>(null)
 
 	// Subscribe to state updates and UI events using the gRPC streaming API
 	useEffect(() => {
@@ -445,6 +440,27 @@ export const ExtensionStateContextProvider: React.FC<{
 			},
 		})
 
+		// Subscribe to theme changes
+		themeSubscriptionRef.current = UiServiceClient.subscribeToTheme(EmptyRequest.create({}), {
+			onResponse: (response) => {
+				if (response.value) {
+					try {
+						const themeData = JSON.parse(response.value)
+						setTheme(convertTextMateToHljs(themeData))
+						console.log("[DEBUG] Received theme update from gRPC stream")
+					} catch (error) {
+						console.error("Error parsing theme data:", error)
+					}
+				}
+			},
+			onError: (error) => {
+				console.error("Error in theme subscription:", error)
+			},
+			onComplete: () => {
+				console.log("Theme subscription completed")
+			},
+		})
+
 		// Still send the webviewDidLaunch message for other initialization
 		vscode.postMessage({ type: "webviewDidLaunch" })
 
@@ -497,6 +513,10 @@ export const ExtensionStateContextProvider: React.FC<{
 				mcpMarketplaceUnsubscribeRef.current()
 				mcpMarketplaceUnsubscribeRef.current = null
 			}
+			if (themeSubscriptionRef.current) {
+				themeSubscriptionRef.current()
+				themeSubscriptionRef.current = null
+			}
 		}
 	}, [])