Fix CodeBlock diffs

webview-ui/src/components/CodeBlock.tsx

@@ -78,7 +78,10 @@ const CodeBlock = ({ code, diff, language, path }: CodeBlockProps) => {
     */
 	const removeLeadingNonAlphanumeric = (path: string): string => path.replace(/^[^a-zA-Z0-9]+/, "")
 
-	const inferredLanguage = useMemo(() => language ?? (path ? getLanguageFromPath(path) : undefined), [path, language])
+	const inferredLanguage = useMemo(
+		() => code && (language ?? (path ? getLanguageFromPath(path) : undefined)),
+		[path, language, code]
+	)
 
 	console.log(inferredLanguage)
 
@@ -126,29 +129,15 @@ const CodeBlock = ({ code, diff, language, path }: CodeBlockProps) => {
 					}}>
 					<SyntaxHighlighter
 						wrapLines={false}
-						language={inferredLanguage}
+						language={diff ? "diff" : inferredLanguage} // "diff" automatically colors changed lines in green/red
 						style={oneDark}
 						customStyle={{
 							margin: 0,
 							padding: "6px 10px",
 							borderRadius: 0,
-                            fontSize: 'var(--vscode-editor-font-size)',
-                            lineHeight: 'var(--vscode-editor-line-height)'
-						}}
-						lineProps={
-							diff != null
-								? (lineNumber) => {
-										const line = diff?.split("\n")?.[lineNumber - 1]
-										let style: React.CSSProperties = { display: "block", width: "100%" }
-										if (line && line[0] === "+") {
-											style.backgroundColor = "var(--vscode-diffEditor-insertedTextBackground)"
-										} else if (line && line[0] === "-") {
-											style.backgroundColor = "var(--vscode-diffEditor-removedTextBackground)"
-										}
-										return { style }
-								  }
-								: undefined
-						}>
+							fontSize: "var(--vscode-editor-font-size)",
+							lineHeight: "var(--vscode-editor-line-height)",
+						}}>
 						{code ?? diff ?? ""}
 					</SyntaxHighlighter>
 				</div>

webview-ui/src/utilities/mockMessages.ts

@@ -126,25 +126,64 @@ export const mockMessages: ClaudeMessage[] = [
 		text: JSON.stringify({
 			tool: "editedExistingFile",
 			path: "/src/components/TodoList.tsx",
-			diff: `@@ -14,6 +14,11 @@
-       }
-     };
+			diff: `+ 
++ 
+  import React, { useState } from "react"
+  
+  interface Todo {
+  	id: number
+  	text: string
+  	completed: boolean
+  }
+  
+  const TodoList: React.FC = () => {
+  	const [todos, setTodos] = useState<Todo[]>([])
+  	const [inputValue, setInputValue] = useState("")
+  
+  	const addTodo = () => {
+  		if (inputValue.trim() !== "") {
+  			setTodos([...todos, { id: Date.now(), text: inputValue, completed: false }])
+  			setInputValue("")
+  		}
+  	}
   
-  +  const deleteTodo = (id: number) => {
-  +    setTodos(todos.filter(todo => todo.id !== id));
-  +  };
-  +
   	const toggleTodo = (id: number) => {
-       setTodos(todos.map(todo =>
-         todo.id === id ? { ...todo, completed: !todo.completed } : todo
-  @@ -35,6 +40,7 @@
-               style={{ textDecoration: todo.completed ? 'line-through' : 'none' }}
-             >
-               {todo.text}
-  +            <button onClick={() => deleteTodo(todo.id)}>Delete</button>
-             </li>
-           ))}
-         </ul>`,
+  		setTodos(todos.map((todo) => (todo.id === id ? { ...todo, completed: !todo.completed } : todo)))
+  	}
+  
+  	return (
+  		<div>
+  			<h1>Todo List</h1>
+  			<input
+  				type="text"
+  				value={inputValue}
+  				onChange={(e) => setInputValue(e.target.value)}
+  				placeholder="Add a new todo"
+  			/>
+  
+- 			<button onClick={addTodo}>Add</button>
+- 			<ul>
+- 				{todos.map((todo) => (
+- 					<li
+- 						key={todo.id}
+- 						onClick={() => toggleTodo(todo.id)}
+- 						style={{ textDecoration: todo.completed ? "line-through" : "none" }}>
+- 						{todo.text}
+- 					</li>
+- 				))}
+- 			</ul>
+- 		</div>
+- 
++ 			
++ 
+  	)
+  }
+  
+  export default TodoList
+  
++ 
++ 		
+`,
 		}),
 	},
 	{