Improve applyDiffs to handle failures robustly (onlook-dev#1663)

* Improve applyDiffs to handle failures robustly

* Handle failing to apply code change

---------

Co-authored-by: Kiet Ho <[email protected]>

Author: devin-ai-integration[bot]

apps/studio/src/lib/editor/engine/chat/code.ts

@@ -3,6 +3,7 @@ import { sendAnalytics } from '@/lib/utils';
 import { CodeBlockProcessor } from '@onlook/ai';
 import { ChatMessageRole, type AssistantChatMessage, type CodeBlock } from '@onlook/models/chat';
 import type { CodeDiff } from '@onlook/models/code';
+import { toast } from '@onlook/ui/use-toast';
 import { makeAutoObservable } from 'mobx';
 import type { ChatManager } from '.';
 import type { EditorEngine } from '..';
@@ -42,7 +43,16 @@ export class ChatCodeManager {
             }
             let content = originalContent;
             for (const block of codeBlocks) {
-                content = await this.processor.applyDiff(content, block.content);
+                const result = await this.processor.applyDiff(content, block.content);
+                if (!result.success) {
+                    console.error('Failed to apply code block', block);
+                    toast({
+                        title: 'Failed to apply code block',
+                        variant: 'destructive',
+                        description: 'Please try again or prompt the AI to fix it.',
+                    });
+                }
+                content = result.text;
             }
 
             const success = await this.writeFileContent(file, content, originalContent);

apps/studio/src/routes/editor/EditPanel/ChatTab/CodeChangeDisplay/CodeModal.tsx

@@ -1,6 +1,6 @@
 import { useEditorEngine } from '@/components/Context';
 import { Button } from '@onlook/ui/button';
-import { Dialog, DialogContent, DialogTrigger } from '@onlook/ui/dialog';
+import { Dialog, DialogContent, DialogTitle, DialogTrigger } from '@onlook/ui/dialog';
 import { Icons } from '@onlook/ui/icons';
 import { Tabs, TabsContent, TabsList, TabsTrigger } from '@onlook/ui/tabs';
 import { useState } from 'react';
@@ -31,6 +31,7 @@ export default function CodeModal({
         <Dialog open={isOpen} onOpenChange={setIsOpen}>
             <DialogTrigger asChild>{children}</DialogTrigger>
             <DialogContent className="min-w-[90vw] h-[80vh]">
+                <DialogTitle className="sr-only">{fileName}</DialogTitle>
                 <Tabs value={selectedTab} onValueChange={(val) => setSelectedTab(val as TabValue)}>
                     <TabsList className="bg-transparent w-full gap-2 justify-start">
                         <TabsTrigger

packages/ai/src/coder/block.ts

@@ -3,17 +3,6 @@ import { FENCE } from '../prompt/format';
 import { flexibleSearchAndReplace } from './search-replace';
 
 export class CodeBlockProcessor {
-    /**
-     * Sequentially applies a list of diffs to the original text
-     */
-    async applyDiffs(originalText: string, diffs: string[]): Promise<string> {
-        let text = originalText;
-        for (const diff of diffs) {
-            text = await this.applyDiff(text, diff);
-        }
-        return text;
-    }
-
     /**
      * Extracts search and replace content from a diff string using the defined fence markers
      */
@@ -58,21 +47,46 @@ export class CodeBlockProcessor {
      * Applies a search/replace diff to the original text with advanced formatting handling
      * Uses multiple strategies and preprocessing options to handle complex replacements
      */
-    async applyDiff(originalText: string, diffText: string): Promise<string> {
+    async applyDiff(
+        originalText: string,
+        diffText: string,
+    ): Promise<{
+        success: boolean;
+        text: string;
+        failures?: Array<{ search: string; error?: string }>;
+    }> {
         const searchReplaces = CodeBlockProcessor.parseDiff(diffText);
         let text = originalText;
+        const failures: Array<{ search: string; error?: string }> = [];
 
         for (const { search, replace } of searchReplaces) {
             const result = await flexibleSearchAndReplace(search, replace, text);
-            if (!result.success) {
-                // Fallback to simple replacement if flexible strategies fail
-                text = text.replace(search, replace);
-            } else if (result.text) {
+            if (result.success && result.text) {
                 text = result.text;
+            } else {
+                // Fallback to simple replacement if flexible strategies fail
+                try {
+                    const newText = text.replace(search, replace);
+                    if (newText !== text) {
+                        text = newText;
+                    } else {
+                        failures.push({ search, error: 'No changes made' });
+                    }
+                } catch (error) {
+                    failures.push({
+                        search,
+                        error: error instanceof Error ? error.message : 'Unknown error',
+                    });
+                    console.warn('Simple replacement failed:', error);
+                }
             }
         }
 
-        return text;
+        return {
+            success: failures.length === 0,
+            text,
+            ...(failures.length > 0 && { failures }),
+        };
     }
 
     /**

packages/ai/test/coder/diffs/diffs.test.ts

@@ -28,7 +28,40 @@ describe('Parse and Apply Code Block Diffs', () => {
         const afterText = await Bun.file(path.resolve(__dirname, './data/single/after.txt')).text();
 
         const result = await coder.applyDiff(beforeText, diffText);
-        expect(result.trim()).toBe(afterText.trim());
+        expect(result.success).toBe(true);
+        expect(result.text.trim()).toBe(afterText.trim());
+        expect(result.failures).toBeUndefined();
+    });
+
+    test('should handle failed replacements', async () => {
+        const diffText = coder.createDiff('non-existent-text', 'replacement');
+        const originalText = 'some sample text';
+
+        const result = await coder.applyDiff(originalText, diffText);
+        expect(result.success).toBe(false);
+        expect(result.text).toBe(originalText);
+        expect(result.failures).toHaveLength(1);
+        expect(result.failures![0]).toEqual({
+            search: 'non-existent-text',
+            error: 'No changes made',
+        });
+    });
+
+    test('should fail when any replacement fails in multiple diffs', async () => {
+        const diffText =
+            coder.createDiff('sample', 'example') +
+            '\n' +
+            coder.createDiff('non-existent', 'replacement');
+        const originalText = 'some sample text';
+
+        const result = await coder.applyDiff(originalText, diffText);
+        expect(result.success).toBe(false);
+        expect(result.text).toBe('some example text');
+        expect(result.failures).toHaveLength(1);
+        expect(result.failures![0]).toEqual({
+            search: 'non-existent',
+            error: 'No changes made',
+        });
     });
 
     test('should create diff correctly', () => {

packages/ai/test/coder/search-replace.test.ts

@@ -99,13 +99,33 @@ describe('CodeBlockProcessor Integration', () => {
         const originalText = '    if (x) {\n        oldFunc();\n    }';
         const diffText = processor.createDiff('    oldFunc();', '    newFunc();');
         const result = await processor.applyDiff(originalText, diffText);
-        expect(result).toBe('    if (x) {\n        newFunc();\n    }');
+        expect(result.success).toBe(true);
+        expect(result.text).toBe('    if (x) {\n        newFunc();\n    }');
+        expect(result.failures).toBeUndefined();
     });
 
     test('should fall back to simple replace if needed', async () => {
         const originalText = 'simple old text';
         const diffText = processor.createDiff('old', 'new');
         const result = await processor.applyDiff(originalText, diffText);
-        expect(result).toBe('simple new text');
+        expect(result.success).toBe(true);
+        expect(result.text).toBe('simple new text');
+        expect(result.failures).toBeUndefined();
+    });
+
+    test('should mark as failed if any replacement fails', async () => {
+        const originalText = 'simple old text';
+        const diffText =
+            processor.createDiff('old', 'new') +
+            '\n' +
+            processor.createDiff('missing', 'replacement');
+        const result = await processor.applyDiff(originalText, diffText);
+        expect(result.success).toBe(false); // Should be false even though one replacement succeeded
+        expect(result.text).toBe('simple new text');
+        expect(result.failures).toHaveLength(1);
+        expect(result.failures![0]).toEqual({
+            search: 'missing',
+            error: 'No changes made',
+        });
     });
 });

packages/ui/src/components/command.tsx

@@ -1,10 +1,10 @@
-import * as React from 'react';
 import type { DialogProps } from '@radix-ui/react-dialog';
 import { MagnifyingGlassIcon } from '@radix-ui/react-icons';
 import { Command as CommandPrimitive } from 'cmdk';
+import * as React from 'react';
 
 import { cn } from '../utils';
-import { Dialog, DialogContent } from './dialog';
+import { Dialog, DialogContent, DialogTitle } from './dialog';
 
 const Command = React.forwardRef<
     React.ElementRef<typeof CommandPrimitive>,
@@ -27,6 +27,7 @@ const CommandDialog = ({ children, ...props }: CommandDialogProps) => {
     return (
         <Dialog {...props}>
             <DialogContent className="overflow-hidden p-0">
+                <DialogTitle className="sr-only">Command</DialogTitle>
                 <Command className="[&_[cmdk-group-heading]]:px-2 [&_[cmdk-group-heading]]:font-medium [&_[cmdk-group-heading]]:text-muted-foreground [&_[cmdk-group]:not([hidden])_~[cmdk-group]]:pt-0 [&_[cmdk-group]]:px-2 [&_[cmdk-input-wrapper]_svg]:h-5 [&_[cmdk-input-wrapper]_svg]:w-5 [&_[cmdk-input]]:h-12 [&_[cmdk-item]]:px-2 [&_[cmdk-item]]:py-3 [&_[cmdk-item]_svg]:h-5 [&_[cmdk-item]_svg]:w-5">
                     {children}
                 </Command>
@@ -133,11 +134,11 @@ CommandShortcut.displayName = 'CommandShortcut';
 export {
     Command,
     CommandDialog,
-    CommandInput,
-    CommandList,
     CommandEmpty,
     CommandGroup,
+    CommandInput,
     CommandItem,
-    CommandShortcut,
+    CommandList,
     CommandSeparator,
+    CommandShortcut,
 };