Add heuristic to suggest fix code output (#45)

Add heuristic to fix code output

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 (

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 && (

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
  */