Finish incomplete dyad write (#475)

Fixes #452 Fixes #456 Fixes #195
2025-06-23 22:53:55 -07:00
parent fa29488b98
commit 3041563809
11 changed files with 1201 additions and 95 deletions
--- a/src/tests/chat_stream_handlers.test.ts
+++ b/src/tests/chat_stream_handlers.test.ts
@@ -6,7 +6,10 @@ import {
  processFullResponseActions,
  getDyadAddDependencyTags,
 } from "../ipc/processors/response_processor";
-import { removeDyadTags } from "../ipc/handlers/chat_stream_handlers";
+import {
+  removeDyadTags,
+  hasUnclosedDyadWrite,
+} from "../ipc/handlers/chat_stream_handlers";
 import fs from "node:fs";
 import git from "isomorphic-git";
 import { db } from "../db";
@@ -1040,7 +1043,7 @@ const component = <Component />;

  it("should handle dyad tags with special characters in content", () => {
    const text = `<dyad-write path="file.js">
-const regex = /<div[^>]*>.*?<\/div>/g;
+const regex = /<div[^>]*>.*?</div>/g;
 const special = "Special chars: @#$%^&*()[]{}|\\";
 </dyad-write>`;
    const result = removeDyadTags(text);
@@ -1059,3 +1062,145 @@ const special = "Special chars: @#$%^&*()[]{}|\\";
    expect(result).toBe("Before  After");
  });
 });
+
+describe("hasUnclosedDyadWrite", () => {
+  it("should return false when there are no dyad-write tags", () => {
+    const text = "This is just regular text without any dyad tags.";
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should return false when dyad-write tag is properly closed", () => {
+    const text = `<dyad-write path="src/file.js">console.log('hello');</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should return true when dyad-write tag is not closed", () => {
+    const text = `<dyad-write path="src/file.js">console.log('hello');`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(true);
+  });
+
+  it("should return false when dyad-write tag with attributes is properly closed", () => {
+    const text = `<dyad-write path="src/file.js" description="A test file">console.log('hello');</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should return true when dyad-write tag with attributes is not closed", () => {
+    const text = `<dyad-write path="src/file.js" description="A test file">console.log('hello');`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(true);
+  });
+
+  it("should return false when there are multiple closed dyad-write tags", () => {
+    const text = `<dyad-write path="src/file1.js">code1</dyad-write>
+    Some text in between
+    <dyad-write path="src/file2.js">code2</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should return true when the last dyad-write tag is unclosed", () => {
+    const text = `<dyad-write path="src/file1.js">code1</dyad-write>
+    Some text in between
+    <dyad-write path="src/file2.js">code2`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(true);
+  });
+
+  it("should return false when first tag is unclosed but last tag is closed", () => {
+    const text = `<dyad-write path="src/file1.js">code1
+    Some text in between
+    <dyad-write path="src/file2.js">code2</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should handle multiline content correctly", () => {
+    const text = `<dyad-write path="src/component.tsx" description="React component">
+import React from 'react';
+
+const Component = () => {
+  return (
+    <div>
+      <h1>Hello World</h1>
+    </div>
+  );
+};
+
+export default Component;
+</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should handle multiline unclosed content correctly", () => {
+    const text = `<dyad-write path="src/component.tsx" description="React component">
+import React from 'react';
+
+const Component = () => {
+  return (
+    <div>
+      <h1>Hello World</h1>
+    </div>
+  );
+};
+
+export default Component;`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(true);
+  });
+
+  it("should handle complex attributes correctly", () => {
+    const text = `<dyad-write path="src/file.js" description="File with quotes and special chars" version="1.0" author="test">
+const message = "Hello 'world'";
+const regex = /<div[^>]*>/g;
+</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should handle text before and after dyad-write tags", () => {
+    const text = `Some text before the tag
+<dyad-write path="src/file.js">console.log('hello');</dyad-write>
+Some text after the tag`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should handle unclosed tag with text after", () => {
+    const text = `Some text before the tag
+<dyad-write path="src/file.js">console.log('hello');
+Some text after the unclosed tag`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(true);
+  });
+
+  it("should handle empty dyad-write tags", () => {
+    const text = `<dyad-write path="src/file.js"></dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should handle unclosed empty dyad-write tags", () => {
+    const text = `<dyad-write path="src/file.js">`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(true);
+  });
+
+  it("should focus on the last opening tag when there are mixed states", () => {
+    const text = `<dyad-write path="src/file1.js">completed content</dyad-write>
+    <dyad-write path="src/file2.js">unclosed content
+    <dyad-write path="src/file3.js">final content</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+
+  it("should handle tags with special characters in attributes", () => {
+    const text = `<dyad-write path="src/file-name_with.special@chars.js" description="File with special chars in path">content</dyad-write>`;
+    const result = hasUnclosedDyadWrite(text);
+    expect(result).toBe(false);
+  });
+});
--- a/src/ipc/handlers/chat_stream_handlers.ts
+++ b/src/ipc/handlers/chat_stream_handlers.ts
@@ -451,41 +451,86 @@ This conversation includes one or more image attachments. When the user uploads
          ];
        }

-        // When calling streamText, the messages need to be properly formatted for mixed content
-        const { fullStream } = streamText({
-          maxTokens: await getMaxTokens(settings.selectedModel),
-          temperature: 0,
-          maxRetries: 2,
-          model: modelClient.model,
-          providerOptions: {
-            "dyad-gateway": getExtraProviderOptions(
-              modelClient.builtinProviderId,
-            ),
-            google: {
-              thinkingConfig: {
-                includeThoughts: true,
-              },
-            } satisfies GoogleGenerativeAIProviderOptions,
-          },
-          system: systemPrompt,
-          messages: chatMessages.filter((m) => m.content),
-          onError: (error: any) => {
-            logger.error("Error streaming text:", error);
-            let errorMessage = (error as any)?.error?.message;
-            const responseBody = error?.error?.responseBody;
-            if (errorMessage && responseBody) {
-              errorMessage += "\n\nDetails: " + responseBody;
-            }
-            const message = errorMessage || JSON.stringify(error);
-            event.sender.send(
-              "chat:response:error",
-              `Sorry, there was an error from the AI: ${message}`,
+        const simpleStreamText = async ({
+          chatMessages,
+        }: {
+          chatMessages: CoreMessage[];
+        }) => {
+          return streamText({
+            maxTokens: await getMaxTokens(settings.selectedModel),
+            temperature: 0,
+            maxRetries: 2,
+            model: modelClient.model,
+            providerOptions: {
+              "dyad-gateway": getExtraProviderOptions(
+                modelClient.builtinProviderId,
+              ),
+              google: {
+                thinkingConfig: {
+                  includeThoughts: true,
+                },
+              } satisfies GoogleGenerativeAIProviderOptions,
+            },
+            system: systemPrompt,
+            messages: chatMessages.filter((m) => m.content),
+            onError: (error: any) => {
+              logger.error("Error streaming text:", error);
+              let errorMessage = (error as any)?.error?.message;
+              const responseBody = error?.error?.responseBody;
+              if (errorMessage && responseBody) {
+                errorMessage += "\n\nDetails: " + responseBody;
+              }
+              const message = errorMessage || JSON.stringify(error);
+              event.sender.send(
+                "chat:response:error",
+                `Sorry, there was an error from the AI: ${message}`,
+              );
+              // Clean up the abort controller
+              activeStreams.delete(req.chatId);
+            },
+            abortSignal: abortController.signal,
+          });
+        };
+
+        const processResponseChunkUpdate = async ({
+          fullResponse,
+        }: {
+          fullResponse: string;
+        }) => {
+          if (
+            fullResponse.includes("$$SUPABASE_CLIENT_CODE$$") &&
+            updatedChat.app?.supabaseProjectId
+          ) {
+            const supabaseClientCode = await getSupabaseClientCode({
+              projectId: updatedChat.app?.supabaseProjectId,
+            });
+            fullResponse = fullResponse.replace(
+              "$$SUPABASE_CLIENT_CODE$$",
+              supabaseClientCode,
            );
-            // Clean up the abort controller
-            activeStreams.delete(req.chatId);
-          },
-          abortSignal: abortController.signal,
-        });
+          }
+          // Store the current partial response
+          partialResponses.set(req.chatId, fullResponse);
+
+          // Update the placeholder assistant message content in the messages array
+          const currentMessages = [...updatedChat.messages];
+          if (
+            currentMessages.length > 0 &&
+            currentMessages[currentMessages.length - 1].role === "assistant"
+          ) {
+            currentMessages[currentMessages.length - 1].content = fullResponse;
+          }
+
+          // Update the assistant message in the database
+          safeSend(event.sender, "chat:response:chunk", {
+            chatId: req.chatId,
+            messages: currentMessages,
+          });
+          return fullResponse;
+        };
+
+        // When calling streamText, the messages need to be properly formatted for mixed content
+        const { fullStream } = await simpleStreamText({ chatMessages });

        // Process the stream as before
        let inThinkingBlock = false;
@@ -520,36 +565,8 @@ This conversation includes one or more image attachments. When the user uploads

            fullResponse += chunk;
            fullResponse = cleanFullResponse(fullResponse);
-
-            if (
-              fullResponse.includes("$$SUPABASE_CLIENT_CODE$$") &&
-              updatedChat.app?.supabaseProjectId
-            ) {
-              const supabaseClientCode = await getSupabaseClientCode({
-                projectId: updatedChat.app?.supabaseProjectId,
-              });
-              fullResponse = fullResponse.replace(
-                "$$SUPABASE_CLIENT_CODE$$",
-                supabaseClientCode,
-              );
-            }
-            // Store the current partial response
-            partialResponses.set(req.chatId, fullResponse);
-
-            // Update the placeholder assistant message content in the messages array
-            const currentMessages = [...updatedChat.messages];
-            if (
-              currentMessages.length > 0 &&
-              currentMessages[currentMessages.length - 1].role === "assistant"
-            ) {
-              currentMessages[currentMessages.length - 1].content =
-                fullResponse;
-            }
-
-            // Update the assistant message in the database
-            safeSend(event.sender, "chat:response:chunk", {
-              chatId: req.chatId,
-              messages: currentMessages,
+            fullResponse = await processResponseChunkUpdate({
+              fullResponse,
            });

            // If the stream was aborted, exit early
@@ -558,6 +575,45 @@ This conversation includes one or more image attachments. When the user uploads
              break;
            }
          }
+
+          if (
+            !abortController.signal.aborted &&
+            settings.selectedChatMode !== "ask" &&
+            hasUnclosedDyadWrite(fullResponse)
+          ) {
+            let continuationAttempts = 0;
+            while (
+              hasUnclosedDyadWrite(fullResponse) &&
+              continuationAttempts < 2 &&
+              !abortController.signal.aborted
+            ) {
+              logger.warn(
+                `Received unclosed dyad-write tag, attempting to continue, attempt #${continuationAttempts + 1}`,
+              );
+              continuationAttempts++;
+
+              const { fullStream: contStream } = await simpleStreamText({
+                // Build messages: replay history then pre-fill assistant with current partial.
+                chatMessages: [
+                  ...chatMessages,
+                  { role: "assistant", content: fullResponse },
+                ],
+              });
+              for await (const part of contStream) {
+                // If the stream was aborted, exit early
+                if (abortController.signal.aborted) {
+                  logger.log(`Stream for chat ${req.chatId} was aborted`);
+                  break;
+                }
+                if (part.type !== "text-delta") continue; // ignore reasoning for continuation
+                fullResponse += part.textDelta;
+                fullResponse = cleanFullResponse(fullResponse);
+                fullResponse = await processResponseChunkUpdate({
+                  fullResponse,
+                });
+              }
+            }
+          }
        } catch (streamError) {
          // Check if this was an abort error
          if (abortController.signal.aborted) {
@@ -832,3 +888,25 @@ export function removeDyadTags(text: string): string {
  const dyadRegex = /<dyad-[^>]*>[\s\S]*?<\/dyad-[^>]*>/g;
  return text.replace(dyadRegex, "").trim();
 }
+
+export function hasUnclosedDyadWrite(text: string): boolean {
+  // Find the last opening dyad-write tag
+  const openRegex = /<dyad-write[^>]*>/g;
+  let lastOpenIndex = -1;
+  let match;
+
+  while ((match = openRegex.exec(text)) !== null) {
+    lastOpenIndex = match.index;
+  }
+
+  // If no opening tag found, there's nothing unclosed
+  if (lastOpenIndex === -1) {
+    return false;
+  }
+
+  // Look for a closing tag after the last opening tag
+  const textAfterLastOpen = text.substring(lastOpenIndex);
+  const hasClosingTag = /<\/dyad-write>/.test(textAfterLastOpen);
+
+  return !hasClosingTag;
+}