[PROTOBUS] Re-enable streaming state, fix the memory leak probably (RooCodeInc#3754)

* Revert "fix: Address memory leak by bypassing subscribeToState gRPC stream for state updates (RooCodeInc#3597)"

This reverts commit 8ab35a5.

* memory leak console boys

* cleanup

---------

Co-authored-by: Andrei Edell <[email protected]>
Co-authored-by: Andrei Eternal <[email protected]>

Author: 3 people

src/core/controller/index.ts

@@ -1178,10 +1178,7 @@ export class Controller {
 
 	async postStateToWebview() {
 		const state = await this.getStateToPostToWebview()
-		// For testing: Bypass gRPC stream and send state directly
-		console.log("[Controller Test Revert] Posting full state via direct 'state' message.")
-		await this.postMessageToWebview({ type: "state", state: state })
-		// await sendStateUpdate(state) // Original line for the GrPC stream
+		await sendStateUpdate(state)
 	}
 
 	async getStateToPostToWebview(): Promise<ExtensionState> {

webview-ui/src/context/ExtensionStateContext.tsx

@@ -207,60 +207,6 @@ export const ExtensionStateContextProvider: React.FC<{
 				}
 				break
 			}
-			case "state": {
-				// Handler for direct state messages
-				if (message.state) {
-					const stateData = message.state as ExtensionState
-					console.log("[Webview Context Test Revert] Received direct 'state' message, updating state.")
-					setState((prevState) => {
-						// Versioning logic for autoApprovalSettings (copied from original onResponse)
-						const incomingVersion = stateData.autoApprovalSettings?.version ?? 1
-						const currentVersion = prevState.autoApprovalSettings?.version ?? 1
-						const shouldUpdateAutoApproval = incomingVersion > currentVersion
-
-						const newState = {
-							...stateData,
-							autoApprovalSettings: shouldUpdateAutoApproval
-								? stateData.autoApprovalSettings
-								: prevState.autoApprovalSettings,
-						}
-
-						// Update welcome screen state based on API configuration (copied from original onResponse)
-						const config = stateData.apiConfiguration
-						const hasKey = config
-							? [
-									config.apiKey,
-									config.openRouterApiKey,
-									config.awsRegion,
-									config.vertexProjectId,
-									config.openAiApiKey,
-									config.ollamaModelId,
-									config.lmStudioModelId,
-									config.liteLlmApiKey,
-									config.geminiApiKey,
-									config.openAiNativeApiKey,
-									config.deepSeekApiKey,
-									config.requestyApiKey,
-									config.togetherApiKey,
-									config.qwenApiKey,
-									config.doubaoApiKey,
-									config.mistralApiKey,
-									config.vsCodeLmModelSelector,
-									config.clineApiKey,
-									config.asksageApiKey,
-									config.xaiApiKey,
-									config.sambanovaApiKey,
-									config.nebiusApiKey,
-								].some((key) => key !== undefined)
-							: false
-
-						setShowWelcome(!hasKey)
-						setDidHydrateState(true)
-						return newState
-					})
-				}
-				break
-			}
 			case "theme": {
 				if (message.text) {
 					setTheme(convertTextMateToHljs(JSON.parse(message.text)))
@@ -326,33 +272,31 @@ export const ExtensionStateContextProvider: React.FC<{
 	const stateSubscriptionRef = useRef<(() => void) | null>(null)
 
 	// Subscribe to state updates using the new gRPC streaming API
-	/* // TEST REVERT: Commenting out gRPC state subscription
 	useEffect(() => {
 		// Set up state subscription
 		stateSubscriptionRef.current = StateServiceClient.subscribeToState(
 			{},
 			{
 				onResponse: (response) => {
-					console.log("[DEBUG] got state update via subscription", response);
 					if (response.stateJson) {
 						try {
-							const stateData = JSON.parse(response.stateJson) as ExtensionState;
-							console.log("[DEBUG] parsed state JSON, updating state");
+							const stateData = JSON.parse(response.stateJson) as ExtensionState
+							console.log("[DEBUG] parsed state JSON, updating state")
 							setState((prevState) => {
 								// Versioning logic for autoApprovalSettings
-								const incomingVersion = stateData.autoApprovalSettings?.version ?? 1;
-								const currentVersion = prevState.autoApprovalSettings?.version ?? 1;
-								const shouldUpdateAutoApproval = incomingVersion > currentVersion;
+								const incomingVersion = stateData.autoApprovalSettings?.version ?? 1
+								const currentVersion = prevState.autoApprovalSettings?.version ?? 1
+								const shouldUpdateAutoApproval = incomingVersion > currentVersion
 
 								const newState = {
 									...stateData,
 									autoApprovalSettings: shouldUpdateAutoApproval
 										? stateData.autoApprovalSettings
 										: prevState.autoApprovalSettings,
-								};
+								}
 
 								// Update welcome screen state based on API configuration
-								const config = stateData.apiConfiguration;
+								const config = stateData.apiConfiguration
 								const hasKey = config
 									? [
 											config.apiKey,
@@ -377,52 +321,41 @@ export const ExtensionStateContextProvider: React.FC<{
 											config.xaiApiKey,
 											config.sambanovaApiKey,
 										].some((key) => key !== undefined)
-									: false;
+									: false
 
-								setShowWelcome(!hasKey);
-								setDidHydrateState(true);
+								setShowWelcome(!hasKey)
+								setDidHydrateState(true)
 
-								console.log("[DEBUG] returning new state in ESC");
+								console.log("[DEBUG] returning new state in ESC")
 
-								return newState;
-							});
+								return newState
+							})
 						} catch (error) {
-							console.error("Error parsing state JSON:", error);
-							console.log("[DEBUG] ERR getting state", error);
+							console.error("Error parsing state JSON:", error)
+							console.log("[DEBUG] ERR getting state", error)
 						}
 					}
-					console.log('[DEBUG] ended "got subscribed state"');
+					console.log('[DEBUG] ended "got subscribed state"')
 				},
 				onError: (error) => {
-					console.error("Error in state subscription:", error);
+					console.error("Error in state subscription:", error)
 				},
 				onComplete: () => {
-					console.log("State subscription completed");
+					console.log("State subscription completed")
 				},
 			},
-		);
+		)
 
 		// Still send the webviewDidLaunch message for other initialization
-		vscode.postMessage({ type: "webviewDidLaunch" });
+		vscode.postMessage({ type: "webviewDidLaunch" })
 
 		// Clean up subscription when component unmounts
 		return () => {
 			if (stateSubscriptionRef.current) {
-				stateSubscriptionRef.current();
-				stateSubscriptionRef.current = null;
+				stateSubscriptionRef.current()
+				stateSubscriptionRef.current = null
 			}
-		};
-	}, []);
-	*/ // END TEST REVERT
-
-	// For the test revert, ensure webviewDidLaunch is still sent if not done by the above useEffect
-	useEffect(() => {
-		// This effect now only sends webviewDidLaunch if the gRPC subscription is commented out.
-		// If the gRPC subscription is active, it sends webviewDidLaunch.
-		// To avoid sending it twice if you uncomment the above, you might add a flag.
-		// For this specific test (gRPC sub commented out), this is fine.
-		console.log("[Webview Context Test Revert] Sending webviewDidLaunch from separate useEffect.")
-		vscode.postMessage({ type: "webviewDidLaunch" })
+		}
 	}, [])
 
 	const contextValue: ExtensionStateContextType = {

webview-ui/src/services/grpc-client-base.ts

@@ -94,7 +94,6 @@ export function createGrpcClient<T extends ProtoService>(service: T): GrpcClient
 							if (message.grpc_response.message) {
 								const responseType = method.responseType
 								const response = responseType.fromJSON(message.grpc_response.message)
-								console.log("[DEBUG] Received streaming response:", message.grpc_response.message)
 								options.onResponse(response)
 							}
 						}
@@ -149,7 +148,6 @@ export function createGrpcClient<T extends ProtoService>(service: T): GrpcClient
 								// Convert JSON back to protobuf message
 								const responseType = method.responseType
 								const response = responseType.fromJSON(message.grpc_response.message)
-								console.log("[DEBUG] grpc-client sending response:", response)
 								resolve(response)
 							}
 						}