Initial streaming refactor

webview-ui/src/components/chat/ChatRow.tsx

@@ -741,10 +741,11 @@ const ProgressIndicator = () => (
 const Markdown = memo(({ markdown }: { markdown?: string }) => {
 	// react-markdown lets us customize elements, so here we're using their example of replacing code blocks with SyntaxHighlighter. However when there are no language matches (` or ``` without a language specifier) then we default to a normal code element for inline code. Code blocks without a language specifier shouldn't be a common occurrence as we prompt Claude to always use a language specifier.
 	// when claude wraps text in thinking tags, he doesnt use line breaks so we need to insert those ourselves to render markdown correctly
-	const parsed = markdown?.replace(/<thinking>([\s\S]*?)<\/thinking>/g, (match, content) => {
-		return content
-		// return `_<thinking>_\n\n${content}\n\n_</thinking>_`
-	})
+	// const parsed = markdown?.replace(/<thinking>([\s\S]*?)<\/thinking>/g, (match, content) => {
+	// 	return content
+	// 	// return `_<thinking>_\n\n${content}\n\n_</thinking>_`
+	// })
+	const parsed = markdown
 	return (
 		<div style={{ wordBreak: "break-word", overflowWrap: "anywhere", marginBottom: -10, marginTop: -10 }}>
 			<ReactMarkdown

webview-ui/src/components/chat/ChatView.tsx

@@ -59,6 +59,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 		if (lastMessage) {
 			switch (lastMessage.type) {
 				case "ask":
+					const isPartial = lastMessage.partial === true
 					switch (lastMessage.ask) {
 						case "api_req_failed":
 							setTextAreaDisabled(true)
@@ -75,16 +76,16 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 							setSecondaryButtonText("Start New Task")
 							break
 						case "followup":
-							setTextAreaDisabled(false)
+							setTextAreaDisabled(isPartial)
 							setClaudeAsk("followup")
-							setEnableButtons(false)
+							setEnableButtons(isPartial)
 							// setPrimaryButtonText(undefined)
 							// setSecondaryButtonText(undefined)
 							break
 						case "tool":
-							setTextAreaDisabled(false)
+							setTextAreaDisabled(isPartial)
 							setClaudeAsk("tool")
-							setEnableButtons(true)
+							setEnableButtons(!isPartial)
 							const tool = JSON.parse(lastMessage.text || "{}") as ClaudeSayTool
 							switch (tool.tool) {
 								case "editedExistingFile":
@@ -99,9 +100,9 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 							}
 							break
 						case "command":
-							setTextAreaDisabled(false)
+							setTextAreaDisabled(isPartial)
 							setClaudeAsk("command")
-							setEnableButtons(true)
+							setEnableButtons(!isPartial)
 							setPrimaryButtonText("Run Command")
 							setSecondaryButtonText("Reject")
 							break
@@ -114,9 +115,9 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 							break
 						case "completion_result":
 							// extension waiting for feedback. but we can just present a new task button
-							setTextAreaDisabled(false)
+							setTextAreaDisabled(isPartial)
 							setClaudeAsk("completion_result")
-							setEnableButtons(true)
+							setEnableButtons(!isPartial)
 							setPrimaryButtonText("Start New Task")
 							setSecondaryButtonText(undefined)
 							break

webview-ui/src/context/ExtensionStateContext.tsx

@@ -4,6 +4,7 @@ import { ExtensionMessage, ExtensionState } from "../../../src/shared/ExtensionM
 import { ApiConfiguration } from "../../../src/shared/api"
 import { vscode } from "../utils/vscode"
 import { convertTextMateToHljs } from "../utils/textMateToHljs"
+import { findLastIndex } from "../../../src/shared/array"
 
 interface ExtensionStateContextType extends ExtensionState {
 	didHydrateState: boolean
@@ -32,29 +33,49 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 
 	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.awsRegion,
-						config.vertexProjectId,
-						config.openAiApiKey,
-						config.ollamaModelId,
-						config.geminiApiKey,
-						config.openAiNativeApiKey,
-				  ].some((key) => key !== undefined)
-				: false
-			setShowWelcome(!hasKey)
-			setDidHydrateState(true)
-		}
-		if (message.type === "theme" && message.text) {
-			setTheme(convertTextMateToHljs(JSON.parse(message.text)))
-		}
-		if (message.type === "workspaceUpdated" && message.filePaths) {
-			setFilePaths(message.filePaths)
+		switch (message.type) {
+			case "state": {
+				setState(message.state!)
+				const config = message.state?.apiConfiguration
+				const hasKey = config
+					? [
+							config.apiKey,
+							config.openRouterApiKey,
+							config.awsRegion,
+							config.vertexProjectId,
+							config.openAiApiKey,
+							config.ollamaModelId,
+							config.geminiApiKey,
+							config.openAiNativeApiKey,
+					  ].some((key) => key !== undefined)
+					: false
+				setShowWelcome(!hasKey)
+				setDidHydrateState(true)
+				break
+			}
+			case "theme": {
+				if (message.text) {
+					setTheme(convertTextMateToHljs(JSON.parse(message.text)))
+				}
+				break
+			}
+			case "workspaceUpdated": {
+				setFilePaths(message.filePaths ?? [])
+				break
+			}
+			case "partialMessage": {
+				const partialMessage = message.partialMessage!
+				setState((prevState) => {
+					// worth noting it will never be possible for a more up-to-date message to be sent here or in normal messages post since the presentAssistantContent function uses lock
+					const lastIndex = findLastIndex(prevState.claudeMessages, (msg) => msg.ts === partialMessage.ts)
+					if (lastIndex !== -1) {
+						const newClaudeMessages = [...prevState.claudeMessages]
+						newClaudeMessages[lastIndex] = partialMessage
+						return { ...prevState, claudeMessages: newClaudeMessages }
+					}
+					return prevState
+				})
+			}
 		}
 	}, [])