Merge pull request #1322 from RooVetGit/cte/fix-extension-state

ExtensionStateContext does not correctly merge state

Author: cte

.changeset/healthy-buckets-attack.md

@@ -0,0 +1,5 @@
+---
+"roo-cline": patch
+---
+
+ExtensionStateContext does not correctly merge state

webview-ui/src/context/ExtensionStateContext.tsx

@@ -69,6 +69,32 @@ export interface ExtensionStateContextType extends ExtensionState {
 
 export const ExtensionStateContext = createContext<ExtensionStateContextType | undefined>(undefined)
 
+export const mergeExtensionState = (prevState: ExtensionState, newState: ExtensionState) => {
+	const {
+		apiConfiguration: prevApiConfiguration,
+		customModePrompts: prevCustomModePrompts,
+		customSupportPrompts: prevCustomSupportPrompts,
+		experiments: prevExperiments,
+		...prevRest
+	} = prevState
+
+	const {
+		apiConfiguration: newApiConfiguration,
+		customModePrompts: newCustomModePrompts,
+		customSupportPrompts: newCustomSupportPrompts,
+		experiments: newExperiments,
+		...newRest
+	} = newState
+
+	const apiConfiguration = { ...prevApiConfiguration, ...newApiConfiguration }
+	const customModePrompts = { ...prevCustomModePrompts, ...newCustomModePrompts }
+	const customSupportPrompts = { ...prevCustomSupportPrompts, ...newCustomSupportPrompts }
+	const experiments = { ...prevExperiments, ...newExperiments }
+	const rest = { ...prevRest, ...newRest }
+
+	return { ...rest, apiConfiguration, customModePrompts, customSupportPrompts, experiments }
+}
+
 export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
 	const [state, setState] = useState<ExtensionState>({
 		version: "",
@@ -123,13 +149,8 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 			switch (message.type) {
 				case "state": {
 					const newState = message.state!
-					setState((prevState) => ({
-						...prevState,
-						...newState,
-					}))
-					const config = newState.apiConfiguration
-					const hasKey = checkExistKey(config)
-					setShowWelcome(!hasKey)
+					setState((prevState) => mergeExtensionState(prevState, newState))
+					setShowWelcome(!checkExistKey(newState.apiConfiguration))
 					setDidHydrateState(true)
 					break
 				}

webview-ui/src/context/__tests__/ExtensionStateContext.test.tsx

@@ -1,6 +1,11 @@
-import React from "react"
+// npx jest webview-ui/src/context/__tests__/ExtensionStateContext.test.tsx
+
 import { render, screen, act } from "@testing-library/react"
-import { ExtensionStateContextProvider, useExtensionState } from "../ExtensionStateContext"
+
+import { ExtensionState } from "../../../../src/shared/ExtensionMessage"
+import { ExtensionStateContextProvider, useExtensionState, mergeExtensionState } from "../ExtensionStateContext"
+import { ExperimentId } from "../../../../src/shared/experiments"
+import { ApiConfiguration } from "../../../../src/shared/api"
 
 // Test component that consumes the context
 const TestComponent = () => {
@@ -63,3 +68,43 @@ describe("ExtensionStateContext", () => {
 		consoleSpy.mockRestore()
 	})
 })
+
+describe("mergeExtensionState", () => {
+	it("should correctly merge extension states", () => {
+		const baseState: ExtensionState = {
+			version: "",
+			mcpEnabled: false,
+			enableMcpServerCreation: false,
+			clineMessages: [],
+			taskHistory: [],
+			shouldShowAnnouncement: false,
+			enableCheckpoints: true,
+			preferredLanguage: "English",
+			writeDelayMs: 1000,
+			requestDelaySeconds: 5,
+			rateLimitSeconds: 0,
+			mode: "default",
+			experiments: {} as Record<ExperimentId, boolean>,
+			customModes: [],
+			maxOpenTabsContext: 20,
+			apiConfiguration: { providerId: "openrouter" } as ApiConfiguration,
+		}
+
+		const prevState: ExtensionState = {
+			...baseState,
+			apiConfiguration: { modelMaxTokens: 1234, modelMaxThinkingTokens: 123 },
+		}
+		const newState: ExtensionState = {
+			...baseState,
+			apiConfiguration: { modelMaxThinkingTokens: 456, modelTemperature: 0.3 },
+		}
+
+		const result = mergeExtensionState(prevState, newState)
+
+		expect(result.apiConfiguration).toEqual({
+			modelMaxTokens: 1234,
+			modelMaxThinkingTokens: 456,
+			modelTemperature: 0.3,
+		})
+	})
+})