From 672bd790fa9321bf0db1dcf2f42bbad0b519f91f Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Tue, 29 Apr 2025 11:34:21 -0700
Subject: [PATCH] Add heuristic to suggest fix code output (#45)

Add heuristic to fix code output
---
 src/components/chat/ChatInput.tsx          | 38 ++++++++++++++++++++++
 src/components/chat/ChatMessage.tsx        | 11 +++++--
 src/components/chat/DyadMarkdownParser.tsx | 11 +++++++
 src/ipc/handlers/proposal_handlers.ts      |  8 +++++
 src/lib/schemas.ts                         |  7 +++-
 5 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/src/components/chat/ChatInput.tsx b/src/components/chat/ChatInput.tsx
index 211be5b..21afc6a 100644
--- a/src/components/chat/ChatInput.tsx
+++ b/src/components/chat/ChatInput.tsx
@@ -385,12 +385,50 @@ function RefactorFileButton({ path }: { path: string }) {
   );
 }
 
+function WriteCodeProperlyButton() {
+  const chatId = useAtomValue(selectedChatIdAtom);
+  const { streamMessage } = useStreamChat();
+  return (
+    <TooltipProvider>
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <Button
+            variant="outline"
+            size="sm"
+            onClick={() => {
+              if (!chatId) {
+                console.error("No chat id found");
+                return;
+              }
+              streamMessage({
+                prompt: `Write the code in the previous message in the correct format using \`<dyad-write>\` tags!`,
+                chatId,
+                redo: false,
+              });
+            }}
+          >
+            Write code properly
+          </Button>
+        </TooltipTrigger>
+        <TooltipContent>
+          <p>
+            Write code properly (useful when AI generates the code in the wrong
+            format)
+          </p>
+        </TooltipContent>
+      </Tooltip>
+    </TooltipProvider>
+  );
+}
+
 function mapActionToButton(action: SuggestedAction) {
   switch (action.id) {
     case "summarize-in-new-chat":
       return <SummarizeInNewChatButton />;
     case "refactor-file":
       return <RefactorFileButton path={action.path} />;
+    case "write-code-properly":
+      return <WriteCodeProperlyButton />;
     default:
       console.error(`Unsupported action: ${action.id}`);
       return (
diff --git a/src/components/chat/ChatMessage.tsx b/src/components/chat/ChatMessage.tsx
index c9aa787..3b363e8 100644
--- a/src/components/chat/ChatMessage.tsx
+++ b/src/components/chat/ChatMessage.tsx
@@ -1,6 +1,9 @@
 import { memo } from "react";
 import type { Message } from "@/ipc/ipc_types";
-import { DyadMarkdownParser } from "./DyadMarkdownParser";
+import {
+  DyadMarkdownParser,
+  VanillaMarkdownParser,
+} from "./DyadMarkdownParser";
 import { motion } from "framer-motion";
 import { useStreamChat } from "@/hooks/useStreamChat";
 import { CheckCircle, XCircle } from "lucide-react";
@@ -64,7 +67,11 @@ const ChatMessage = ({ message }: ChatMessageProps) => {
             className="prose dark:prose-invert prose-headings:mb-2 prose-p:my-1 prose-pre:my-0 max-w-none"
             suppressHydrationWarning
           >
-            <DyadMarkdownParser content={message.content} />
+            {message.role === "assistant" ? (
+              <DyadMarkdownParser content={message.content} />
+            ) : (
+              <VanillaMarkdownParser content={message.content} />
+            )}
           </div>
         )}
         {message.approvalState && (
diff --git a/src/components/chat/DyadMarkdownParser.tsx b/src/components/chat/DyadMarkdownParser.tsx
index 39d4ff6..86114fc 100644
--- a/src/components/chat/DyadMarkdownParser.tsx
+++ b/src/components/chat/DyadMarkdownParser.tsx
@@ -29,6 +29,17 @@ type ContentPiece =
   | { type: "markdown"; content: string }
   | { type: "custom-tag"; tagInfo: CustomTagInfo };
 
+export const VanillaMarkdownParser = ({ content }: { content: string }) => {
+  return (
+    <ReactMarkdown
+      rehypePlugins={[rehypeRaw]}
+      components={{ code: CodeHighlight } as any}
+    >
+      {content}
+    </ReactMarkdown>
+  );
+};
+
 /**
  * Custom component to parse markdown content with Dyad-specific tags
  */
diff --git a/src/ipc/handlers/proposal_handlers.ts b/src/ipc/handlers/proposal_handlers.ts
index 7db7570..e527e1b 100644
--- a/src/ipc/handlers/proposal_handlers.ts
+++ b/src/ipc/handlers/proposal_handlers.ts
@@ -239,6 +239,14 @@ const getProposalHandler = async (
             path: refactorTarget.path,
           });
         }
+        if (
+          writeTags.length === 0 &&
+          latestAssistantMessage.content.includes("```")
+        ) {
+          actions.push({
+            id: "write-code-properly",
+          });
+        }
       }
 
       // Get all chat messages to calculate token usage
diff --git a/src/lib/schemas.ts b/src/lib/schemas.ts
index b5d6d32..5271f04 100644
--- a/src/lib/schemas.ts
+++ b/src/lib/schemas.ts
@@ -159,7 +159,8 @@ export interface CodeProposal {
 export type SuggestedAction =
   | RestartAppAction
   | SummarizeInNewChatAction
-  | RefactorFileAction;
+  | RefactorFileAction
+  | WriteCodeProperlyAction;
 
 export interface RestartAppAction {
   id: "restart-app";
@@ -169,6 +170,10 @@ export interface SummarizeInNewChatAction {
   id: "summarize-in-new-chat";
 }
 
+export interface WriteCodeProperlyAction {
+  id: "write-code-properly";
+}
+
 export interface RefactorFileAction {
   id: "refactor-file";
   path: string;