From 26d7dfa8f7aa3e36e9112ba872ce9cf64afa0ac5 Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Thu, 29 Aug 2024 04:31:52 -0400
Subject: [PATCH] Refactor showWelcome

---
 webview-ui/src/App.tsx                         | 18 +++++-------------
 .../src/context/ExtensionStateContext.tsx      | 16 ++++++++++++----
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/webview-ui/src/App.tsx b/webview-ui/src/App.tsx
index 64e284d..f2cbfb1 100644
--- a/webview-ui/src/App.tsx
+++ b/webview-ui/src/App.tsx
@@ -10,26 +10,14 @@ import { ExtensionStateContextProvider, useExtensionState } from "./context/Exte
 import { vscode } from "./utils/vscode"
 
 const AppContent = () => {
-	const { apiConfiguration, shouldShowAnnouncement } = useExtensionState()
+	const { didHydrateState, showWelcome, apiConfiguration, shouldShowAnnouncement } = useExtensionState()
 	const [showSettings, setShowSettings] = useState(false)
 	const [showHistory, setShowHistory] = useState(false)
-	const [showWelcome, setShowWelcome] = useState<boolean>(false)
 	const [showAnnouncement, setShowAnnouncement] = useState(false)
 
 	const handleMessage = useCallback((e: MessageEvent) => {
 		const message: ExtensionMessage = e.data
 		switch (message.type) {
-			case "state":
-				let hasKey = false
-				const config = message.state?.apiConfiguration
-				if (config) {
-					const { apiKey, openRouterApiKey, awsAccessKey, vertexProjectId } = config
-					hasKey = [apiKey, openRouterApiKey, awsAccessKey, vertexProjectId].some((key) => key !== undefined)
-				} else {
-					hasKey = false
-				}
-				setShowWelcome(!hasKey)
-				break
 			case "action":
 				switch (message.action!) {
 					case "settingsButtonTapped":
@@ -62,6 +50,10 @@ const AppContent = () => {
 		}
 	}, [shouldShowAnnouncement])
 
+	if (!didHydrateState) {
+		return null
+	}
+
 	return (
 		<>
 			{showWelcome ? (
diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx
index 60fda79..e346229 100644
--- a/webview-ui/src/context/ExtensionStateContext.tsx
+++ b/webview-ui/src/context/ExtensionStateContext.tsx
@@ -5,6 +5,8 @@ import { ApiConfiguration } from "../../../src/shared/api"
 import { vscode } from "../utils/vscode"
 
 interface ExtensionStateContextType extends ExtensionState {
+	didHydrateState: boolean
+	showWelcome: boolean
 	setApiConfiguration: (config: ApiConfiguration) => void
 	setMaxRequestsPerTask: (value?: number) => void
 	setCustomInstructions: (value?: string) => void
@@ -22,11 +24,19 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 		shouldShowAnnouncement: false,
 	})
 	const [didHydrateState, setDidHydrateState] = useState(false)
+	const [showWelcome, setShowWelcome] = useState(false)
 
 	const handleMessage = useCallback((event: MessageEvent) => {
 		const message: ExtensionMessage = event.data
 		if (message.type === "state" && message.state) {
 			setState(message.state)
+			const config = message.state?.apiConfiguration
+			const hasKey = config
+				? [config.apiKey, config.openRouterApiKey, config.awsAccessKey, config.vertexProjectId].some(
+						(key) => key !== undefined
+				  )
+				: false
+			setShowWelcome(!hasKey)
 			setDidHydrateState(true)
 		}
 	}, [])
@@ -39,6 +49,8 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 
 	const contextValue: ExtensionStateContextType = {
 		...state,
+		didHydrateState,
+		showWelcome,
 		setApiConfiguration: (value) => setState((prevState) => ({ ...prevState, apiConfiguration: value })),
 		setMaxRequestsPerTask: (value) => setState((prevState) => ({ ...prevState, maxRequestsPerTask: value })),
 		setCustomInstructions: (value) => setState((prevState) => ({ ...prevState, customInstructions: value })),
@@ -46,10 +58,6 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 		setShowAnnouncement: (value) => setState((prevState) => ({ ...prevState, shouldShowAnnouncement: value })),
 	}
 
-	if (!didHydrateState) {
-		return null
-	}
-
 	return <ExtensionStateContext.Provider value={contextValue}>{children}</ExtensionStateContext.Provider>
 }