import { v4 as uuidv4 } from "uuid";
import { ipcMain, IpcMainInvokeEvent } from "electron";
import { createTypedHandler } from "./base";
import { chatContracts } from "../types/chat";
import {
  ModelMessage,
  TextPart,
  ImagePart,
  streamText,
  ToolSet,
  TextStreamPart,
  stepCountIs,
  hasToolCall,
  type ToolExecutionOptions,
} from "ai";

import { db } from "../../db";
import { chats, messages } from "../../db/schema";
import { and, eq, isNull } from "drizzle-orm";
import type { SmartContextMode } from "../../lib/schemas";
import {
  constructSystemPrompt,
  readAiRules,
} from "../../prompts/system_prompt";
import { getThemePromptById } from "../utils/theme_utils";
import {
  getSupabaseAvailableSystemPrompt,
  SUPABASE_NOT_AVAILABLE_SYSTEM_PROMPT,
} from "../../prompts/supabase_prompt";
import { getDyadAppPath } from "../../paths/paths";
import { readSettings } from "../../main/settings";
import type { ChatResponseEnd, ChatStreamParams } from "@/ipc/types";
import {
  CodebaseFile,
  extractCodebase,
  readFileWithCache,
} from "../../utils/codebase";
import {
  dryRunSearchReplace,
  processFullResponseActions,
} from "../processors/response_processor";
import { streamTestResponse } from "./testing_chat_handlers";
import { getTestResponse } from "./testing_chat_handlers";
import { getModelClient, ModelClient } from "../utils/get_model_client";
import log from "electron-log";
import { sendTelemetryEvent } from "../utils/telemetry";
import {
  getSupabaseContext,
  getSupabaseClientCode,
} from "../../supabase_admin/supabase_context";
import { SUMMARIZE_CHAT_SYSTEM_PROMPT } from "../../prompts/summarize_chat_system_prompt";
import { SECURITY_REVIEW_SYSTEM_PROMPT } from "../../prompts/security_review_prompt";
import fs from "node:fs";
import * as path from "path";
import * as crypto from "crypto";
import { readFile, writeFile } from "fs/promises";
import { getMaxTokens, getTemperature } from "../utils/token_utils";
import { MAX_CHAT_TURNS_IN_CONTEXT } from "@/constants/settings_constants";
import { validateChatContext } from "../utils/context_paths_utils";
import { getProviderOptions, getAiHeaders } from "../utils/provider_options";
import { mcpServers } from "../../db/schema";
import { requireMcpToolConsent } from "../utils/mcp_consent";

import { handleLocalAgentStream } from "../../pro/main/ipc/handlers/local_agent/local_agent_handler";

import { safeSend } from "../utils/safe_sender";
import { cleanFullResponse } from "../utils/cleanFullResponse";
import { generateProblemReport } from "../processors/tsc";
import { createProblemFixPrompt } from "@/shared/problem_prompt";
import { AsyncVirtualFileSystem } from "../../../shared/VirtualFilesystem";
import { escapeXmlAttr, escapeXmlContent } from "../../../shared/xmlEscape";
import {
  getDyadAddDependencyTags,
  getDyadWriteTags,
  getDyadDeleteTags,
  getDyadRenameTags,
} from "../utils/dyad_tag_parser";
import { fileExists } from "../utils/file_utils";
import { extractMentionedAppsCodebases } from "../utils/mention_apps";
import { parseAppMentions } from "@/shared/parse_mention_apps";
import { prompts as promptsTable } from "../../db/schema";
import { inArray } from "drizzle-orm";
import { replacePromptReference } from "../utils/replacePromptReference";
import { replaceSlashSkillReference } from "../utils/replaceSlashSkillReference";
import { parsePlanFile, validatePlanId } from "./planUtils";
import { ensureDyadGitignored } from "./gitignoreUtils";
import { DYAD_MEDIA_DIR_NAME } from "../utils/media_path_utils";
import { mcpManager } from "../utils/mcp_manager";
import z from "zod";
import {
  isBasicAgentMode,
  isSupabaseConnected,
  isTurboEditsV2Enabled,
} from "@/lib/schemas";
import {
  getFreeAgentQuotaStatus,
  markMessageAsUsingFreeAgentQuota,
  unmarkMessageAsUsingFreeAgentQuota,
} from "./free_agent_quota_handlers";
import { AI_STREAMING_ERROR_MESSAGE_PREFIX } from "@/shared/texts";
import { getCurrentCommitHash } from "../utils/git_utils";
import {
  processChatMessagesWithVersionedFiles as getVersionedFiles,
  VersionedFiles,
} from "../utils/versioned_codebase_context";
import { getAiMessagesJsonIfWithinLimit } from "../utils/ai_messages_utils";

type AsyncIterableStream<T> = AsyncIterable<T> & ReadableStream<T>;

const logger = log.scope("chat_stream_handlers");

// Track active streams for cancellation
const activeStreams = new Map<number, AbortController>();

// Track partial responses for cancelled streams
const partialResponses = new Map<number, string>();

// Common helper functions
const TEXT_FILE_EXTENSIONS = [
  ".md",
  ".txt",
  ".json",
  ".csv",
  ".js",
  ".ts",
  ".html",
  ".css",
];

async function isTextFile(filePath: string): Promise<boolean> {
  const ext = path.extname(filePath).toLowerCase();
  return TEXT_FILE_EXTENSIONS.includes(ext);
}

// Use escapeXmlAttr from shared/xmlEscape for XML escaping

// Safely parse an MCP tool key that combines server and tool names.
// We split on the LAST occurrence of "__" to avoid ambiguity if either
// side contains "__" as part of its sanitized name.
function parseMcpToolKey(toolKey: string): {
  serverName: string;
  toolName: string;
} {
  const separator = "__";
  const lastIndex = toolKey.lastIndexOf(separator);
  if (lastIndex === -1) {
    return { serverName: "", toolName: toolKey };
  }
  const serverName = toolKey.slice(0, lastIndex);
  const toolName = toolKey.slice(lastIndex + separator.length);
  return { serverName, toolName };
}

// Helper function to process stream chunks
async function processStreamChunks({
  fullStream,
  fullResponse,
  abortController,
  chatId,
  processResponseChunkUpdate,
}: {
  fullStream: AsyncIterableStream<TextStreamPart<ToolSet>>;
  fullResponse: string;
  abortController: AbortController;
  chatId: number;
  processResponseChunkUpdate: (params: {
    fullResponse: string;
  }) => Promise<string>;
}): Promise<{ fullResponse: string; incrementalResponse: string }> {
  let incrementalResponse = "";
  let inThinkingBlock = false;

  for await (const part of fullStream) {
    let chunk = "";
    if (
      inThinkingBlock &&
      !["reasoning-delta", "reasoning-end", "reasoning-start"].includes(
        part.type,
      )
    ) {
      chunk = "</think>";
      inThinkingBlock = false;
    }
    if (part.type === "text-delta") {
      chunk += part.text;
    } else if (part.type === "reasoning-delta") {
      if (!inThinkingBlock) {
        chunk = "<think>";
        inThinkingBlock = true;
      }

      chunk += escapeDyadTags(part.text);
    } else if (part.type === "tool-call") {
      const { serverName, toolName } = parseMcpToolKey(part.toolName);
      const content = escapeDyadTags(JSON.stringify(part.input));
      chunk = `<dyad-mcp-tool-call server="${serverName}" tool="${toolName}">\n${content}\n</dyad-mcp-tool-call>\n`;
    } else if (part.type === "tool-result") {
      const { serverName, toolName } = parseMcpToolKey(part.toolName);
      const content = escapeDyadTags(part.output);
      chunk = `<dyad-mcp-tool-result server="${serverName}" tool="${toolName}">\n${content}\n</dyad-mcp-tool-result>\n`;
    }

    if (!chunk) {
      continue;
    }

    fullResponse += chunk;
    incrementalResponse += chunk;
    fullResponse = cleanFullResponse(fullResponse);
    fullResponse = await processResponseChunkUpdate({
      fullResponse,
    });

    // If the stream was aborted, exit early
    if (abortController.signal.aborted) {
      logger.log(`Stream for chat ${chatId} was aborted`);
      break;
    }
  }

  return { fullResponse, incrementalResponse };
}

export function registerChatStreamHandlers() {
  ipcMain.handle("chat:stream", async (event, req: ChatStreamParams) => {
    let attachmentPaths: string[] = [];
    try {
      let dyadRequestId: string | undefined;
      // Create an AbortController for this stream
      const abortController = new AbortController();
      activeStreams.set(req.chatId, abortController);

      // Notify renderer that stream is starting
      safeSend(event.sender, "chat:stream:start", { chatId: req.chatId });

      // Get the chat to check for existing messages
      const chat = await db.query.chats.findFirst({
        where: eq(chats.id, req.chatId),
        with: {
          messages: {
            orderBy: (messages, { asc }) => [asc(messages.createdAt)],
          },
          app: true, // Include app information
        },
      });

      if (!chat) {
        throw new Error(`Chat not found: ${req.chatId}`);
      }

      // Handle redo option: remove the most recent messages if needed
      if (req.redo) {
        // Get the most recent messages
        const chatMessages = [...chat.messages];

        // Find the most recent user message
        let lastUserMessageIndex = chatMessages.length - 1;
        while (
          lastUserMessageIndex >= 0 &&
          chatMessages[lastUserMessageIndex].role !== "user"
        ) {
          lastUserMessageIndex--;
        }

        if (lastUserMessageIndex >= 0) {
          // Delete the user message
          await db
            .delete(messages)
            .where(eq(messages.id, chatMessages[lastUserMessageIndex].id));

          // If there's an assistant message after the user message, delete it too
          if (
            lastUserMessageIndex < chatMessages.length - 1 &&
            chatMessages[lastUserMessageIndex + 1].role === "assistant"
          ) {
            await db
              .delete(messages)
              .where(
                eq(messages.id, chatMessages[lastUserMessageIndex + 1].id),
              );
          }
        }
      }

      // Process attachments if any
      let attachmentInfo = "";
      // Display-only attachment info uses <dyad-attachment> tags for inline rendering
      let displayAttachmentInfo = "";

      if (req.attachments && req.attachments.length > 0) {
        attachmentInfo = "\n\nAttachments:\n";

        // Create persistent .dyad/media directory for this app
        const appPath = getDyadAppPath(chat.app.path);
        const mediaDir = path.join(appPath, DYAD_MEDIA_DIR_NAME);
        if (!fs.existsSync(mediaDir)) {
          fs.mkdirSync(mediaDir, { recursive: true });
        }
        await ensureDyadGitignored(appPath);

        for (let i = 0; i < req.attachments.length; i++) {
          const attachment = req.attachments[i];
          // Generate a unique filename (include index to avoid collisions
          // when multiple attachments share the same name within the same ms)
          const hash = crypto
            .createHash("md5")
            .update(attachment.name + Date.now() + i)
            .digest("hex");
          const fileExtension = path.extname(attachment.name);
          const filename = `${hash}${fileExtension}`;

          // Extract the base64 data (remove the data:mime/type;base64, prefix)
          const base64Data = attachment.data.split(";base64,").pop() || "";
          const fileBuffer = Buffer.from(base64Data, "base64");

          // Save to .dyad/media dir
          const persistentPath = path.join(mediaDir, filename);
          await writeFile(persistentPath, fileBuffer);
          attachmentPaths.push(persistentPath);

          // Build dyad-media:// URL for display
          // Use a fixed hostname to avoid URL hostname normalization (lowercasing)
          // Encode path segments so special characters (spaces, #, ?, %) don't
          // break URL parsing. The protocol handler already decodeURIComponent's.
          const mediaUrl = `dyad-media://media/${encodeURIComponent(chat.app.path)}/.dyad/media/${encodeURIComponent(filename)}`;

          // Build display tag for inline rendering (escape attribute values)
          displayAttachmentInfo += `\n<dyad-attachment name="${escapeXmlAttr(attachment.name)}" type="${escapeXmlAttr(attachment.type)}" url="${escapeXmlAttr(mediaUrl)}" path="${escapeXmlAttr(persistentPath)}" attachment-type="${escapeXmlAttr(attachment.attachmentType)}"></dyad-attachment>\n`;

          if (attachment.attachmentType === "upload-to-codebase") {
            // Provide the .dyad/media path so the AI can copy it into the codebase
            attachmentInfo += `\n\nFile to upload to codebase: "${attachment.name}" (path: ${persistentPath})\nUse the copy_file tool (or <dyad-copy> tag) to copy this file into the codebase at the appropriate location.\n`;
          } else {
            // For chat-context, provide file info for reference (no path to avoid auto-copying)
            attachmentInfo += `- ${attachment.name} (${attachment.type})\n`;
            // If it's a text-based file, try to include the content
            if (await isTextFile(persistentPath)) {
              try {
                attachmentInfo += `<dyad-text-attachment filename="${escapeXmlAttr(attachment.name)}" type="${escapeXmlAttr(attachment.type)}" path="${escapeXmlAttr(persistentPath)}">
                </dyad-text-attachment>
                \n\n`;
              } catch (err) {
                logger.error(`Error reading file content: ${err}`);
              }
            }
          }
        }
      }

      // Build the full AI prompt (with .dyad/media paths and copy_file instructions)
      let userPrompt = req.prompt + (attachmentInfo ? attachmentInfo : "");
      // Build the display prompt (with <dyad-attachment> tags for inline rendering)
      // This separates what the user sees from what the AI receives.
      let displayUserPrompt: string | undefined;
      if (displayAttachmentInfo) {
        displayUserPrompt = req.prompt + displayAttachmentInfo;
      }
      // Inline referenced prompt contents for mentions like @prompt:<id>
      try {
        const matches = Array.from(userPrompt.matchAll(/@prompt:(\d+)/g));
        if (matches.length > 0) {
          const ids = Array.from(new Set(matches.map((m) => Number(m[1]))));
          const referenced = await db
            .select()
            .from(promptsTable)
            .where(inArray(promptsTable.id, ids));
          if (referenced.length > 0) {
            const promptsMap: Record<number, string> = {};
            for (const p of referenced) {
              promptsMap[p.id] = p.content;
            }
            userPrompt = replacePromptReference(userPrompt, promptsMap);
          }
        }
      } catch (e) {
        logger.error("Failed to inline referenced prompts:", e);
      }

      // Expand /slug skill references (e.g. /webapp-testing) to prompt content
      try {
        const slashSkillPattern = /(?:^|\s)\/([a-zA-Z0-9-]+)(?=\s|$)/;
        if (slashSkillPattern.test(userPrompt)) {
          const allPrompts = db.select().from(promptsTable).all();
          const promptsBySlug: Record<string, string> = {};
          for (const p of allPrompts) {
            if (p.slug && !promptsBySlug[p.slug]) {
              promptsBySlug[p.slug] = p.content;
            }
          }
          userPrompt = replaceSlashSkillReference(userPrompt, promptsBySlug);
        }
      } catch (e) {
        logger.error("Failed to expand slash skill references:", e);
      }

      // Expand /implement-plan= into full implementation prompt
      // Keep the original short form for display in the UI; the expanded
      // content is only injected into the AI message history.
      let implementPlanDisplayPrompt: string | undefined;
      const implementPlanMatch = userPrompt.match(/^\/implement-plan=(.+)$/);
      if (implementPlanMatch) {
        try {
          implementPlanDisplayPrompt = userPrompt;
          const planSlug = implementPlanMatch[1];
          validatePlanId(planSlug);
          const appPath = getDyadAppPath(chat.app.path);
          const planFilePath = path.join(
            appPath,
            ".dyad",
            "plans",
            `${planSlug}.md`,
          );
          const raw = await fs.promises.readFile(planFilePath, "utf-8");
          const { meta, content } = parsePlanFile(raw);

          const planPath = `.dyad/plans/${planSlug}.md`;

          userPrompt = `Please implement the following plan:

## ${meta.title || "Implementation Plan"}

${content}

Start implementing this plan now. Follow the steps outlined and create/modify the necessary files.
You may update the plan at \`${planPath}\` to mark your progress.`;
        } catch (e) {
          implementPlanDisplayPrompt = undefined;
          logger.error("Failed to expand /implement-plan= prompt:", e);
        }
      }

      const componentsToProcess = req.selectedComponents || [];

      if (componentsToProcess.length > 0) {
        userPrompt += "\n\nSelected components:\n";

        for (const component of componentsToProcess) {
          let componentSnippet = "[component snippet not available]";
          try {
            const componentFileContent = await readFile(
              path.join(getDyadAppPath(chat.app.path), component.relativePath),
              "utf8",
            );
            const lines = componentFileContent.split(/\r?\n/);
            const selectedIndex = component.lineNumber - 1;

            // Let's get one line before and three after for context.
            const startIndex = Math.max(0, selectedIndex - 1);
            const endIndex = Math.min(lines.length, selectedIndex + 4);

            const snippetLines = lines.slice(startIndex, endIndex);
            const selectedLineInSnippetIndex = selectedIndex - startIndex;

            if (snippetLines[selectedLineInSnippetIndex]) {
              snippetLines[selectedLineInSnippetIndex] =
                `${snippetLines[selectedLineInSnippetIndex]} // <-- EDIT HERE`;
            }

            componentSnippet = snippetLines.join("\n");
          } catch (err) {
            logger.error(
              `Error reading selected component file content: ${err}`,
            );
          }

          userPrompt += `\n${componentsToProcess.length > 1 ? `${componentsToProcess.indexOf(component) + 1}. ` : ""}Component: ${component.name} (file: ${component.relativePath})

Snippet:
\`\`\`
${componentSnippet}
\`\`\`
`;
        }
      }

      const [insertedUserMessage] = await db
        .insert(messages)
        .values({
          chatId: req.chatId,
          role: "user",
          content:
            implementPlanDisplayPrompt ?? displayUserPrompt ?? userPrompt,
        })
        .returning({ id: messages.id });
      const userMessageId = insertedUserMessage.id;
      const settings = readSettings();
      // Only Dyad Pro requests have request ids.
      if (settings.enableDyadPro) {
        // Generate requestId early so it can be saved with the message
        dyadRequestId = uuidv4();
      }

      // Add a placeholder assistant message immediately
      const [placeholderAssistantMessage] = await db
        .insert(messages)
        .values({
          chatId: req.chatId,
          role: "assistant",
          content: "", // Start with empty content
          requestId: dyadRequestId,
          model: settings.selectedModel.name,
          sourceCommitHash: await getCurrentCommitHash({
            path: getDyadAppPath(chat.app.path),
          }),
        })
        .returning();

      // Fetch updated chat data after possible deletions and additions
      const updatedChat = await db.query.chats.findFirst({
        where: eq(chats.id, req.chatId),
        with: {
          messages: {
            orderBy: (messages, { asc }) => [asc(messages.createdAt)],
          },
          app: true, // Include app information
        },
      });

      if (!updatedChat) {
        throw new Error(`Chat not found: ${req.chatId}`);
      }

      // Send the messages right away so that the loading state is shown for the message.
      safeSend(event.sender, "chat:response:chunk", {
        chatId: req.chatId,
        messages: updatedChat.messages,
      });

      let fullResponse = "";
      let maxTokensUsed: number | undefined;

      // Check if this is a test prompt
      const testResponse = getTestResponse(req.prompt);

      if (testResponse) {
        // For test prompts, use the dedicated function
        fullResponse = await streamTestResponse(
          event,
          req.chatId,
          testResponse,
          abortController,
          updatedChat,
        );
      } else {
        // Normal AI processing for non-test prompts
        const { modelClient, isEngineEnabled, isSmartContextEnabled } =
          await getModelClient(settings.selectedModel, settings);

        const appPath = getDyadAppPath(updatedChat.app.path);
        // When we don't have smart context enabled, we
        // only include the selected components' files for codebase context.
        //
        // If we have selected components and smart context is enabled,
        // we handle this specially below.
        const chatContext =
          req.selectedComponents &&
          req.selectedComponents.length > 0 &&
          !isSmartContextEnabled
            ? {
                contextPaths: req.selectedComponents.map((component) => ({
                  globPath: component.relativePath,
                })),
                smartContextAutoIncludes: [],
              }
            : validateChatContext(updatedChat.app.chatContext);

        // Extract codebase for current app
        const { formattedOutput: codebaseInfo, files } = await extractCodebase({
          appPath,
          chatContext,
        });

        // For smart context and selected components, we will mark the selected components' files as focused.
        // This means that we don't do the regular smart context handling, but we'll allow fetching
        // additional files through <dyad-read> as needed.
        if (
          isSmartContextEnabled &&
          req.selectedComponents &&
          req.selectedComponents.length > 0
        ) {
          const selectedPaths = new Set(
            req.selectedComponents.map((component) => component.relativePath),
          );
          for (const file of files) {
            if (selectedPaths.has(file.path)) {
              file.focused = true;
            }
          }
        }

        // Parse app mentions from the prompt
        const mentionedAppNames = parseAppMentions(req.prompt);

        // Extract codebases for mentioned apps
        const mentionedAppsCodebases = await extractMentionedAppsCodebases(
          mentionedAppNames,
          updatedChat.app.id, // Exclude current app
        );
        const willUseLocalAgentStream =
          (settings.selectedChatMode === "local-agent" ||
            settings.selectedChatMode === "ask") &&
          !mentionedAppsCodebases.length;

        const isDeepContextEnabled =
          isEngineEnabled &&
          settings.enableProSmartFilesContextMode &&
          // Anything besides balanced will use deep context.
          settings.proSmartContextOption !== "balanced" &&
          mentionedAppsCodebases.length === 0;
        logger.log(`isDeepContextEnabled: ${isDeepContextEnabled}`);

        // Combine current app codebase with mentioned apps' codebases
        let otherAppsCodebaseInfo = "";
        if (mentionedAppsCodebases.length > 0) {
          const mentionedAppsSection = mentionedAppsCodebases
            .map(
              ({ appName, codebaseInfo }) =>
                `\n\n=== Referenced App: ${appName} ===\n${codebaseInfo}`,
            )
            .join("");

          otherAppsCodebaseInfo = mentionedAppsSection;

          logger.log(
            `Added ${mentionedAppsCodebases.length} mentioned app codebases`,
          );
        }

        logger.log(`Extracted codebase information from ${appPath}`);
        logger.log(
          "codebaseInfo: length",
          codebaseInfo.length,
          "estimated tokens",
          codebaseInfo.length / 4,
        );

        // Prepare message history for the AI
        const messageHistory = updatedChat.messages.map((message) => ({
          role: message.role as "user" | "assistant" | "system",
          content: message.content,
          sourceCommitHash: message.sourceCommitHash,
          commitHash: message.commitHash,
        }));

        // The DB stores display-friendly versions (short /implement-plan= form
        // or clean <dyad-attachment> tags). Replace the last user message with the
        // full AI prompt so the model receives expanded plan content or attachment paths.
        if (implementPlanDisplayPrompt || displayUserPrompt) {
          for (let i = messageHistory.length - 1; i >= 0; i--) {
            if (messageHistory[i].role === "user") {
              messageHistory[i] = {
                ...messageHistory[i],
                content: userPrompt,
              };
              break;
            }
          }
        }

        // For Dyad Pro + Deep Context, we set to 200 chat turns (+1)
        // this is to enable more cache hits. Practically, users should
        // rarely go over this limit because they will hit the model's
        // context window limit.
        //
        // Limit chat history based on maxChatTurnsInContext setting
        // We add 1 because the current prompt counts as a turn.
        const maxChatTurns = isDeepContextEnabled
          ? 201
          : (settings.maxChatTurnsInContext || MAX_CHAT_TURNS_IN_CONTEXT) + 1;

        // If we need to limit the context, we take only the most recent turns
        let limitedMessageHistory = messageHistory;
        if (messageHistory.length > maxChatTurns * 2) {
          // Each turn is a user + assistant pair
          // Calculate how many messages to keep (maxChatTurns * 2)
          let recentMessages = messageHistory
            .filter((msg) => msg.role !== "system")
            .slice(-maxChatTurns * 2);

          // Ensure the first message is a user message
          if (recentMessages.length > 0 && recentMessages[0].role !== "user") {
            // Find the first user message
            const firstUserIndex = recentMessages.findIndex(
              (msg) => msg.role === "user",
            );
            if (firstUserIndex > 0) {
              // Drop assistant messages before the first user message
              recentMessages = recentMessages.slice(firstUserIndex);
            } else if (firstUserIndex === -1) {
              logger.warn(
                "No user messages found in recent history, set recent messages to empty",
              );
              recentMessages = [];
            }
          }

          limitedMessageHistory = [...recentMessages];

          logger.log(
            `Limiting chat history from ${messageHistory.length} to ${limitedMessageHistory.length} messages (max ${maxChatTurns} turns)`,
          );
        }

        const aiRules = await readAiRules(getDyadAppPath(updatedChat.app.path));

        // Get theme prompt for the app (null themeId means "no theme")
        const themePrompt = await getThemePromptById(updatedChat.app.themeId);
        logger.log(
          `Theme for app ${updatedChat.app.id}: ${updatedChat.app.themeId ?? "none"}, prompt length: ${themePrompt.length} chars`,
        );

        // Migration on read converts "agent" to "build", so no need to check for it here
        let systemPrompt = constructSystemPrompt({
          aiRules,
          chatMode: settings.selectedChatMode,
          enableTurboEditsV2: isTurboEditsV2Enabled(settings),
          themePrompt,
          basicAgentMode: isBasicAgentMode(settings),
        });

        // Add information about mentioned apps if any
        if (otherAppsCodebaseInfo) {
          const mentionedAppsList = mentionedAppsCodebases
            .map(({ appName }) => appName)
            .join(", ");

          systemPrompt += `\n\n# Referenced Apps\nThe user has mentioned the following apps in their prompt: ${mentionedAppsList}. Their codebases have been included in the context for your reference. When referring to these apps, you can understand their structure and code to provide better assistance, however you should NOT edit the files in these referenced apps. The referenced apps are NOT part of the current app and are READ-ONLY.`;
        }

        const isSecurityReviewIntent =
          req.prompt.startsWith("/security-review");
        if (isSecurityReviewIntent) {
          systemPrompt = SECURITY_REVIEW_SYSTEM_PROMPT;
          try {
            const appPath = getDyadAppPath(updatedChat.app.path);
            const rulesPath = path.join(appPath, "SECURITY_RULES.md");
            let securityRules = "";

            await fs.promises.access(rulesPath);
            securityRules = await fs.promises.readFile(rulesPath, "utf8");

            if (securityRules && securityRules.trim().length > 0) {
              systemPrompt +=
                "\n\n# Project-specific security rules:\n" + securityRules;
            }
          } catch (error) {
            // Best-effort: if reading rules fails, continue without them
            logger.info("Failed to read security rules", error);
          }
        }

        if (
          updatedChat.app?.supabaseProjectId &&
          isSupabaseConnected(settings)
        ) {
          const supabaseClientCode = await getSupabaseClientCode({
            projectId: updatedChat.app.supabaseProjectId,
            organizationSlug: updatedChat.app.supabaseOrganizationSlug ?? null,
          });
          systemPrompt +=
            "\n\n" +
            getSupabaseAvailableSystemPrompt(supabaseClientCode) +
            "\n\n" +
            // For local agent, we will explicitly fetch the database context when needed.
            (settings.selectedChatMode === "local-agent"
              ? ""
              : await getSupabaseContext({
                  supabaseProjectId: updatedChat.app.supabaseProjectId,
                  organizationSlug:
                    updatedChat.app.supabaseOrganizationSlug ?? null,
                }));
        } else if (
          // Neon projects don't need Supabase.
          !updatedChat.app?.neonProjectId &&
          // In local agent mode, we will suggest supabase as part of the add-integration tool
          settings.selectedChatMode !== "local-agent" &&
          // If in security review mode, we don't need to mention supabase is available.
          !isSecurityReviewIntent
        ) {
          systemPrompt += "\n\n" + SUPABASE_NOT_AVAILABLE_SYSTEM_PROMPT;
        }
        const isSummarizeIntent = req.prompt.startsWith(
          "Summarize from chat-id=",
        );
        if (isSummarizeIntent) {
          systemPrompt = SUMMARIZE_CHAT_SYSTEM_PROMPT;
        }

        // Update the system prompt for images if there are image attachments
        const hasImageAttachments =
          req.attachments &&
          req.attachments.some((attachment) =>
            attachment.type.startsWith("image/"),
          );

        const hasUploadedAttachments =
          req.attachments &&
          req.attachments.some(
            (attachment) => attachment.attachmentType === "upload-to-codebase",
          );
        // If there's mixed attachments (e.g. some upload to codebase attachments and some upload images as chat context attachemnts)
        // we will just include the file upload system prompt, otherwise the AI gets confused and doesn't reliably
        // print out the dyad-write tags.
        // Usually, AI models will want to use the image as reference to generate code (e.g. UI mockups) anyways, so
        // it's not that critical to include the image analysis instructions.
        const isAskMode = settings.selectedChatMode === "ask";
        if (hasUploadedAttachments) {
          if (willUseLocalAgentStream && !isAskMode) {
            systemPrompt += `

When files are attached for upload to the codebase, use the \`copy_file\` tool to copy them from their path into the project.

Example:
\`\`\`
copy_file(from=".dyad/media/abc123.png", to="src/assets/logo.png", description="Copy uploaded image into project")
\`\`\`

The file paths are provided in the attachment information above.
`;
          } else if (!isAskMode) {
            systemPrompt += `

When files are attached for upload to the codebase, copy them into the project using this format:

<dyad-copy from=".dyad/media/abc123.png" to="src/assets/logo.png" description="Copy uploaded file"></dyad-copy>

The file paths are provided in the attachment information above.
`;
          }
        } else if (hasImageAttachments) {
          systemPrompt += `

# Image Analysis Instructions
This conversation includes one or more image attachments. When the user uploads images:
1. If the user explicitly asks for analysis, description, or information about the image, please analyze the image content.
2. Describe what you see in the image if asked.
3. You can use images as references when the user has coding or design-related questions.
4. For diagrams or wireframes, try to understand the content and structure shown.
5. For screenshots of code or errors, try to identify the issue or explain the code.
`;
        }

        const codebasePrefix = isEngineEnabled
          ? // No codebase prefix if engine is set, we will take of it there.
            []
          : ([
              {
                role: "user",
                content: createCodebasePrompt(codebaseInfo),
              },
              {
                role: "assistant",
                content: "OK, got it. I'm ready to help",
              },
            ] as const);

        // If engine is enabled, we will send the other apps codebase info to the engine
        // and process it with smart context.
        const otherCodebasePrefix =
          otherAppsCodebaseInfo && !isEngineEnabled
            ? ([
                {
                  role: "user",
                  content: createOtherAppsCodebasePrompt(otherAppsCodebaseInfo),
                },
                {
                  role: "assistant",
                  content: "OK.",
                },
              ] as const)
            : [];

        const limitedHistoryChatMessages = limitedMessageHistory.map((msg) => ({
          role: msg.role as "user" | "assistant" | "system",
          // Why remove thinking tags?
          // Thinking tags are generally not critical for the context
          // and eats up extra tokens.
          content:
            settings.selectedChatMode === "ask"
              ? removeDyadTags(removeNonEssentialTags(msg.content))
              : removeNonEssentialTags(msg.content),
          providerOptions: {
            "dyad-engine": {
              sourceCommitHash: msg.sourceCommitHash,
              commitHash: msg.commitHash,
            },
          },
        }));

        let chatMessages: ModelMessage[] = [
          ...codebasePrefix,
          ...otherCodebasePrefix,
          ...limitedHistoryChatMessages,
        ];

        // Check if the last message should include attachments
        if (chatMessages.length >= 2) {
          const lastUserIndex = chatMessages.length - 2;
          const lastUserMessage = chatMessages[lastUserIndex];
          if (lastUserMessage.role === "user") {
            if (attachmentPaths.length > 0) {
              // Replace the last message with one that includes attachments
              chatMessages[lastUserIndex] = await prepareMessageWithAttachments(
                lastUserMessage,
                attachmentPaths,
              );
            }
            // Save aiMessagesJson for modes that use handleLocalAgentStream
            // (which reads from DB and needs structured image content)

            if (willUseLocalAgentStream) {
              // Insert into DB (with size guard)
              const userAiMessagesJson = getAiMessagesJsonIfWithinLimit([
                chatMessages[lastUserIndex],
              ]);
              if (userAiMessagesJson) {
                await db
                  .update(messages)
                  .set({ aiMessagesJson: userAiMessagesJson })
                  .where(eq(messages.id, userMessageId));
              }
            }
          }
        } else {
          logger.warn(
            "Unexpected number of chat messages:",
            chatMessages.length,
          );
        }

        if (isSummarizeIntent) {
          const previousChat = await db.query.chats.findFirst({
            where: eq(chats.id, parseInt(req.prompt.split("=")[1])),
            with: {
              messages: {
                orderBy: (messages, { asc }) => [asc(messages.createdAt)],
              },
            },
          });
          chatMessages = [
            {
              role: "user",
              content:
                "Summarize the following chat: " +
                formatMessagesForSummary(previousChat?.messages ?? []),
            } satisfies ModelMessage,
          ];
        }
        const simpleStreamText = async ({
          chatMessages,
          modelClient,
          tools,
          systemPromptOverride = systemPrompt,
          dyadDisableFiles = false,
          files,
        }: {
          chatMessages: ModelMessage[];
          modelClient: ModelClient;
          files: CodebaseFile[];
          tools?: ToolSet;
          systemPromptOverride?: string;
          dyadDisableFiles?: boolean;
        }) => {
          if (isEngineEnabled) {
            logger.log(
              "sending AI request to engine with request id:",
              dyadRequestId,
            );
          } else {
            logger.log("sending AI request");
          }
          let versionedFiles: VersionedFiles | undefined;
          if (isDeepContextEnabled) {
            versionedFiles = await getVersionedFiles({
              files,
              chatMessages,
              appPath,
            });
          }
          const smartContextMode: SmartContextMode = isDeepContextEnabled
            ? "deep"
            : "balanced";
          const providerOptions = getProviderOptions({
            dyadAppId: updatedChat.app.id,
            dyadRequestId,
            dyadDisableFiles,
            smartContextMode,
            files,
            versionedFiles,
            mentionedAppsCodebases,
            builtinProviderId: modelClient.builtinProviderId,
            settings,
          });

          const streamResult = streamText({
            headers: getAiHeaders({
              builtinProviderId: modelClient.builtinProviderId,
            }),
            maxOutputTokens: await getMaxTokens(settings.selectedModel),
            temperature: await getTemperature(settings.selectedModel),
            maxRetries: 2,
            model: modelClient.model,
            stopWhen: [stepCountIs(20), hasToolCall("edit-code")],
            providerOptions,
            system: systemPromptOverride,
            tools,
            messages: chatMessages.filter((m) => m.content),
            onFinish: (response) => {
              const totalTokens = response.usage?.totalTokens;

              if (typeof totalTokens === "number") {
                // We use the highest total tokens used (we are *not* accumulating)
                // since we're trying to figure it out if we're near the context limit.
                maxTokensUsed = Math.max(maxTokensUsed ?? 0, totalTokens);

                // Persist the aggregated token usage on the placeholder assistant message
                void db
                  .update(messages)
                  .set({ maxTokensUsed: maxTokensUsed })
                  .where(eq(messages.id, placeholderAssistantMessage.id))
                  .catch((error) => {
                    logger.error(
                      "Failed to save total tokens for assistant message",
                      error,
                    );
                  });

                logger.log(
                  `Total tokens used (aggregated for message ${placeholderAssistantMessage.id}): ${maxTokensUsed}`,
                );
              } else {
                logger.log("Total tokens used: unknown");
              }
            },
            onError: (error: any) => {
              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);
              const requestIdPrefix = isEngineEnabled
                ? `[Request ID: ${dyadRequestId}] `
                : "";
              logger.error(
                `AI stream text error for request: ${requestIdPrefix} errorMessage=${errorMessage} error=`,
                error,
              );
              event.sender.send("chat:response:error", {
                chatId: req.chatId,
                error: `${AI_STREAMING_ERROR_MESSAGE_PREFIX}${requestIdPrefix}${message}`,
              });
              // Clean up the abort controller
              activeStreams.delete(req.chatId);
            },
            abortSignal: abortController.signal,
          });
          return {
            fullStream: streamResult.fullStream,
            usage: streamResult.usage,
          };
        };

        let lastDbSaveAt = 0;

        const processResponseChunkUpdate = async ({
          fullResponse,
        }: {
          fullResponse: string;
        }) => {
          // Store the current partial response
          partialResponses.set(req.chatId, fullResponse);
          // Save to DB (in case user is switching chats during the stream)
          const now = Date.now();
          if (now - lastDbSaveAt >= 150) {
            await db
              .update(messages)
              .set({ content: fullResponse })
              .where(eq(messages.id, placeholderAssistantMessage.id));

            lastDbSaveAt = now;
          }

          // 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;
        };

        // Handle ask mode: use local-agent in read-only mode
        // This gives users access to code reading tools while in ask mode
        // Ask mode does not consume free agent quota
        if (
          settings.selectedChatMode === "ask" &&
          !mentionedAppsCodebases.length
        ) {
          // Reconstruct system prompt for local-agent read-only mode
          const readOnlySystemPrompt = constructSystemPrompt({
            aiRules,
            chatMode: "local-agent",
            enableTurboEditsV2: false,
            themePrompt,
            readOnly: true,
          });

          // Return value indicates success/failure for quota tracking.
          // Ask mode doesn't consume quota, but we still capture it for
          // consistent error handling.
          const streamSuccess = await handleLocalAgentStream(
            event,
            req,
            abortController,
            {
              placeholderMessageId: placeholderAssistantMessage.id,
              // Note: this is using the read-only system prompt rather than the
              // regular system prompt which gets overrides for special intents
              // like summarize chat, security review, etc.
              //
              // This is OK because those intents should always happen in a new chat
              // and new chats will default to non-ask modes.
              systemPrompt: readOnlySystemPrompt,
              dyadRequestId: dyadRequestId ?? "[no-request-id]",
              readOnly: true,
              messageOverride: isSummarizeIntent ? chatMessages : undefined,
            },
          );
          if (!streamSuccess) {
            logger.warn(
              "Ask mode local agent stream did not complete successfully",
            );
          }
          return;
        }

        // Handle plan mode: use local-agent with plan tools only
        // Plan mode is for requirements gathering and creating implementation plans
        if (
          settings.selectedChatMode === "plan" &&
          !mentionedAppsCodebases.length
        ) {
          // Reconstruct system prompt for plan mode
          const planModeSystemPrompt = constructSystemPrompt({
            aiRules,
            chatMode: "plan",
            enableTurboEditsV2: false,
            themePrompt,
          });

          await handleLocalAgentStream(event, req, abortController, {
            placeholderMessageId: placeholderAssistantMessage.id,
            systemPrompt: planModeSystemPrompt,
            dyadRequestId: dyadRequestId ?? "[no-request-id]",
            planModeOnly: true,
            messageOverride: isSummarizeIntent ? chatMessages : undefined,
          });
          return;
        }

        // Handle local-agent mode (Agent v2)
        // Mentioned apps can't be handled by the local agent (defer to balanced smart context
        // in build mode)
        if (
          settings.selectedChatMode === "local-agent" &&
          !mentionedAppsCodebases.length
        ) {
          // Check quota for Basic Agent mode (non-Pro users)
          const isBasicAgentModeRequest = isBasicAgentMode(settings);
          if (isBasicAgentModeRequest) {
            const quotaStatus = await getFreeAgentQuotaStatus();
            if (quotaStatus.isQuotaExceeded) {
              safeSend(event.sender, "chat:response:error", {
                chatId: req.chatId,
                error: JSON.stringify({
                  type: "FREE_AGENT_QUOTA_EXCEEDED",
                  hoursUntilReset: quotaStatus.hoursUntilReset,
                  resetTime: quotaStatus.resetTime,
                }),
              });
              return;
            }
          }

          // Mark the user message as using quota BEFORE starting the stream
          // to prevent race conditions with parallel requests
          if (isBasicAgentModeRequest && userMessageId) {
            await markMessageAsUsingFreeAgentQuota(userMessageId);
          }

          let streamSuccess = false;
          try {
            streamSuccess = await handleLocalAgentStream(
              event,
              req,
              abortController,
              {
                placeholderMessageId: placeholderAssistantMessage.id,
                systemPrompt,
                dyadRequestId: dyadRequestId ?? "[no-request-id]",
                messageOverride: isSummarizeIntent ? chatMessages : undefined,
              },
            );
          } finally {
            // If the stream failed, was aborted, or threw, refund the quota
            if (isBasicAgentModeRequest && userMessageId && !streamSuccess) {
              await unmarkMessageAsUsingFreeAgentQuota(userMessageId);
            }
          }

          return;
        }

        // Use MCP agent code path if:
        // 1. The enableMcpServersForBuildMode experiment is on AND
        // 2. Mode is "build" AND there are enabled MCP servers
        if (
          settings.enableMcpServersForBuildMode &&
          settings.selectedChatMode === "build"
        ) {
          const tools = await getMcpTools(event);
          const hasEnabledMcpServers = Object.keys(tools).length > 0;

          // Only run MCP agent path if build mode has enabled MCP servers
          if (hasEnabledMcpServers) {
            const { fullStream } = await simpleStreamText({
              chatMessages: limitedHistoryChatMessages,
              modelClient,
              tools: {
                ...tools,
                "generate-code": {
                  description:
                    "ALWAYS use this tool whenever generating or editing code for the codebase.",
                  inputSchema: z.object({}),
                  execute: async () => "",
                },
              },
              systemPromptOverride: constructSystemPrompt({
                aiRules: await readAiRules(
                  getDyadAppPath(updatedChat.app.path),
                ),
                chatMode: "build",
                enableTurboEditsV2: false,
              }),
              files: files,
              dyadDisableFiles: true,
            });

            const result = await processStreamChunks({
              fullStream,
              fullResponse,
              abortController,
              chatId: req.chatId,
              processResponseChunkUpdate,
            });
            fullResponse = result.fullResponse;
            chatMessages.push({
              role: "assistant",
              content: fullResponse,
            });
            chatMessages.push({
              role: "user",
              content: "OK.",
            });
          }
        }

        // When calling streamText, the messages need to be properly formatted for mixed content
        const { fullStream } = await simpleStreamText({
          chatMessages,
          modelClient,
          files: files,
        });

        // Process the stream as before
        try {
          const result = await processStreamChunks({
            fullStream,
            fullResponse,
            abortController,
            chatId: req.chatId,
            processResponseChunkUpdate,
          });
          fullResponse = result.fullResponse;

          if (
            settings.selectedChatMode !== "ask" &&
            isTurboEditsV2Enabled(settings)
          ) {
            let issues = await dryRunSearchReplace({
              fullResponse,
              appPath: getDyadAppPath(updatedChat.app.path),
            });
            sendTelemetryEvent("search_replace:fix", {
              attemptNumber: 0,
              success: issues.length === 0,
              issueCount: issues.length,
              errors: issues.map((i) => ({
                filePath: i.filePath,
                error: i.error,
              })),
            });

            let searchReplaceFixAttempts = 0;
            const originalFullResponse = fullResponse;
            const previousAttempts: ModelMessage[] = [];
            while (
              issues.length > 0 &&
              searchReplaceFixAttempts < 2 &&
              !abortController.signal.aborted
            ) {
              logger.warn(
                `Detected search-replace issues (attempt #${searchReplaceFixAttempts + 1}): ${issues.map((i) => i.error).join(", ")}`,
              );
              const formattedSearchReplaceIssues = issues
                .map(({ filePath, error }) => {
                  return `File path: ${filePath}\nError: ${error}`;
                })
                .join("\n\n");

              fullResponse += `<dyad-output type="warning" message="Could not apply Turbo Edits properly for some of the files; re-generating code...">${formattedSearchReplaceIssues}</dyad-output>`;
              await processResponseChunkUpdate({
                fullResponse,
              });

              logger.info(
                `Attempting to fix search-replace issues, attempt #${searchReplaceFixAttempts + 1}`,
              );

              const fixSearchReplacePrompt =
                searchReplaceFixAttempts === 0
                  ? `There was an issue with the following \`dyad-search-replace\` tags. Make sure you use \`dyad-read\` to read the latest version of the file and then trying to do search & replace again.`
                  : `There was an issue with the following \`dyad-search-replace\` tags. Please fix the errors by generating the code changes using \`dyad-write\` tags instead.`;
              searchReplaceFixAttempts++;
              const userPrompt = {
                role: "user",
                content: `${fixSearchReplacePrompt}
                
${formattedSearchReplaceIssues}`,
              } as const;

              const { fullStream: fixSearchReplaceStream } =
                await simpleStreamText({
                  // Build messages: reuse chat history and original full response, then ask to fix search-replace issues.
                  chatMessages: [
                    ...chatMessages,
                    { role: "assistant", content: originalFullResponse },
                    ...previousAttempts,
                    userPrompt,
                  ],
                  modelClient,
                  files: files,
                });
              previousAttempts.push(userPrompt);
              const result = await processStreamChunks({
                fullStream: fixSearchReplaceStream,
                fullResponse,
                abortController,
                chatId: req.chatId,
                processResponseChunkUpdate,
              });
              fullResponse = result.fullResponse;
              previousAttempts.push({
                role: "assistant",
                content: removeNonEssentialTags(result.incrementalResponse),
              });

              // Re-check for issues after the fix attempt
              issues = await dryRunSearchReplace({
                fullResponse: result.incrementalResponse,
                appPath: getDyadAppPath(updatedChat.app.path),
              });

              sendTelemetryEvent("search_replace:fix", {
                attemptNumber: searchReplaceFixAttempts,
                success: issues.length === 0,
                issueCount: issues.length,
                errors: issues.map((i) => ({
                  filePath: i.filePath,
                  error: i.error,
                })),
              });
            }
          }

          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 },
                ],
                modelClient,
                files: files,
              });
              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.text;
                fullResponse = cleanFullResponse(fullResponse);
                fullResponse = await processResponseChunkUpdate({
                  fullResponse,
                });
              }
            }
          }
          const addDependencies = getDyadAddDependencyTags(fullResponse);
          if (
            !abortController.signal.aborted &&
            // If there are dependencies, we don't want to auto-fix problems
            // because there's going to be type errors since the packages aren't
            // installed yet.
            addDependencies.length === 0 &&
            settings.enableAutoFixProblems &&
            settings.selectedChatMode !== "ask"
          ) {
            try {
              // IF auto-fix is enabled
              let problemReport = await generateProblemReport({
                fullResponse,
                appPath: getDyadAppPath(updatedChat.app.path),
              });

              let autoFixAttempts = 0;
              const originalFullResponse = fullResponse;
              const previousAttempts: ModelMessage[] = [];
              while (
                problemReport.problems.length > 0 &&
                autoFixAttempts < 2 &&
                !abortController.signal.aborted
              ) {
                fullResponse += `<dyad-problem-report summary="${problemReport.problems.length} problems">
${problemReport.problems
  .map(
    (problem) =>
      `<problem file="${escapeXmlAttr(problem.file)}" line="${problem.line}" column="${problem.column}" code="${problem.code}">${escapeXmlContent(problem.message)}</problem>`,
  )
  .join("\n")}
</dyad-problem-report>`;

                logger.info(
                  `Attempting to auto-fix problems, attempt #${autoFixAttempts + 1}`,
                );
                autoFixAttempts++;
                const problemFixPrompt = createProblemFixPrompt(problemReport);

                const virtualFileSystem = new AsyncVirtualFileSystem(
                  getDyadAppPath(updatedChat.app.path),
                  {
                    fileExists: (fileName: string) => fileExists(fileName),
                    readFile: (fileName: string) => readFileWithCache(fileName),
                  },
                );
                const writeTags = getDyadWriteTags(fullResponse);
                const renameTags = getDyadRenameTags(fullResponse);
                const deletePaths = getDyadDeleteTags(fullResponse);
                virtualFileSystem.applyResponseChanges({
                  deletePaths,
                  renameTags,
                  writeTags,
                });

                const { formattedOutput: codebaseInfo, files } =
                  await extractCodebase({
                    appPath,
                    chatContext,
                    virtualFileSystem,
                  });
                const { modelClient } = await getModelClient(
                  settings.selectedModel,
                  settings,
                );

                const { fullStream } = await simpleStreamText({
                  modelClient,
                  files: files,
                  chatMessages: [
                    ...chatMessages.map((msg, index) => {
                      if (
                        index === 0 &&
                        msg.role === "user" &&
                        typeof msg.content === "string" &&
                        msg.content.startsWith(CODEBASE_PROMPT_PREFIX)
                      ) {
                        return {
                          role: "user",
                          content: createCodebasePrompt(codebaseInfo),
                        } as const;
                      }
                      return msg;
                    }),
                    {
                      role: "assistant",
                      content: removeNonEssentialTags(originalFullResponse),
                    },
                    ...previousAttempts,
                    { role: "user", content: problemFixPrompt },
                  ],
                });
                previousAttempts.push({
                  role: "user",
                  content: problemFixPrompt,
                });
                const result = await processStreamChunks({
                  fullStream,
                  fullResponse,
                  abortController,
                  chatId: req.chatId,
                  processResponseChunkUpdate,
                });
                fullResponse = result.fullResponse;
                previousAttempts.push({
                  role: "assistant",
                  content: removeNonEssentialTags(result.incrementalResponse),
                });

                problemReport = await generateProblemReport({
                  fullResponse,
                  appPath: getDyadAppPath(updatedChat.app.path),
                });
              }
            } catch (error) {
              logger.error(
                "Error generating problem report or auto-fixing:",
                settings.enableAutoFixProblems,
                error,
              );
            }
          }
        } catch (streamError) {
          // Check if this was an abort error
          if (abortController.signal.aborted) {
            const chatId = req.chatId;
            const partialResponse = partialResponses.get(req.chatId);
            // If we have a partial response, save it to the database
            if (partialResponse) {
              try {
                // Update the placeholder assistant message with the partial content and cancellation note
                await db
                  .update(messages)
                  .set({
                    content: `${partialResponse}

[Response cancelled by user]`,
                  })
                  .where(eq(messages.id, placeholderAssistantMessage.id));

                logger.log(
                  `Updated cancelled response for placeholder message ${placeholderAssistantMessage.id} in chat ${chatId}`,
                );
                partialResponses.delete(req.chatId);
              } catch (error) {
                logger.error(
                  `Error saving partial response for chat ${chatId}:`,
                  error,
                );
              }
            }
            return req.chatId;
          }
          throw streamError;
        }
      }

      // Only save the response and process it if we weren't aborted
      if (!abortController.signal.aborted && fullResponse) {
        // Scrape from: <dyad-chat-summary>Renaming profile file</dyad-chat-title>
        const chatTitle = fullResponse.match(
          /<dyad-chat-summary>(.*?)<\/dyad-chat-summary>/,
        );
        if (chatTitle) {
          await db
            .update(chats)
            .set({ title: chatTitle[1] })
            .where(and(eq(chats.id, req.chatId), isNull(chats.title)));
        }
        const chatSummary = chatTitle?.[1];

        // Update the placeholder assistant message with the full response
        await db
          .update(messages)
          .set({ content: fullResponse })
          .where(eq(messages.id, placeholderAssistantMessage.id));
        const settings = readSettings();
        if (
          settings.autoApproveChanges &&
          settings.selectedChatMode !== "ask"
        ) {
          const status = await processFullResponseActions(
            fullResponse,
            req.chatId,
            {
              chatSummary,
              messageId: placeholderAssistantMessage.id,
            }, // Use placeholder ID
          );

          const chat = await db.query.chats.findFirst({
            where: eq(chats.id, req.chatId),
            with: {
              messages: {
                orderBy: (messages, { asc }) => [asc(messages.createdAt)],
              },
            },
          });

          safeSend(event.sender, "chat:response:chunk", {
            chatId: req.chatId,
            messages: chat!.messages,
          });

          if (status.error) {
            safeSend(event.sender, "chat:response:error", {
              chatId: req.chatId,
              error: `Sorry, there was an error applying the AI's changes: ${status.error}`,
            });
          }

          // Signal that the stream has completed
          safeSend(event.sender, "chat:response:end", {
            chatId: req.chatId,
            updatedFiles: status.updatedFiles ?? false,
            extraFiles: status.extraFiles,
            extraFilesError: status.extraFilesError,
            chatSummary,
          } satisfies ChatResponseEnd);
        } else {
          safeSend(event.sender, "chat:response:end", {
            chatId: req.chatId,
            updatedFiles: false,
            chatSummary,
          } satisfies ChatResponseEnd);
        }
      }

      // Return the chat ID for backwards compatibility
      return req.chatId;
    } catch (error) {
      logger.error("Error calling LLM:", error);
      safeSend(event.sender, "chat:response:error", {
        chatId: req.chatId,
        error: `Sorry, there was an error processing your request: ${error}`,
      });

      return "error";
    } finally {
      // Clean up the abort controller
      activeStreams.delete(req.chatId);

      // Notify renderer that stream has ended
      safeSend(event.sender, "chat:stream:end", { chatId: req.chatId });
    }
  });

  // Handler to cancel an ongoing stream
  createTypedHandler(chatContracts.cancelStream, async (event, chatId) => {
    const abortController = activeStreams.get(chatId);

    if (abortController) {
      // Abort the stream
      abortController.abort();
      activeStreams.delete(chatId);
      logger.log(`Aborted stream for chat ${chatId}`);
    } else {
      logger.warn(`No active stream found for chat ${chatId}`);
    }

    // Send the end event to the renderer with wasCancelled flag
    safeSend(event.sender, "chat:response:end", {
      chatId,
      updatedFiles: false,
      wasCancelled: true,
    } satisfies ChatResponseEnd);

    // Also emit stream:end so cleanup listeners (e.g., pending agent consents) fire
    safeSend(event.sender, "chat:stream:end", { chatId });

    return true;
  });
}

export function formatMessagesForSummary(
  messages: { role: string; content: string | undefined }[],
) {
  if (messages.length <= 8) {
    // If we have 8 or fewer messages, include all of them
    return messages
      .map((m) => `<message role="${m.role}">${m.content}</message>`)
      .join("\n");
  }

  // Take first 2 messages and last 6 messages
  const firstMessages = messages.slice(0, 2);
  const lastMessages = messages.slice(-6);

  // Combine them with an indicator of skipped messages
  const combinedMessages = [
    ...firstMessages,
    {
      role: "system",
      content: `[... ${messages.length - 8} messages omitted ...]`,
    },
    ...lastMessages,
  ];

  return combinedMessages
    .map((m) => `<message role="${m.role}">${m.content}</message>`)
    .join("\n");
}

// Helper function to replace text attachment placeholders with full content
async function replaceTextAttachmentWithContent(
  text: string,
  filePath: string,
  fileName: string,
): Promise<string> {
  try {
    if (await isTextFile(filePath)) {
      // Read the full content
      const fullContent = await readFile(filePath, "utf-8");

      // Replace the placeholder tag with the full content.
      // The path attribute in the tag is XML-escaped (via escapeXmlAttr), so we
      // must also XML-escape the path before regex-escaping to ensure a match.
      const xmlEscapedPath = escapeXmlAttr(filePath);
      const escapedPath = xmlEscapedPath.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
      const tagPattern = new RegExp(
        `<dyad-text-attachment filename="[^"]*" type="[^"]*" path="${escapedPath}">\\s*<\\/dyad-text-attachment>`,
        "g",
      );

      const replacedText = text.replace(
        tagPattern,
        `Full content of ${fileName}:\n\`\`\`\n${fullContent}\n\`\`\``,
      );

      logger.log(
        `Replaced text attachment content for: ${fileName} - length before: ${text.length} - length after: ${replacedText.length}`,
      );
      return replacedText;
    }
    return text;
  } catch (error) {
    logger.error(`Error processing text file: ${error}`);
    return text;
  }
}

// Helper function to convert traditional message to one with proper image attachments
async function prepareMessageWithAttachments(
  message: ModelMessage,
  attachmentPaths: string[],
): Promise<ModelMessage> {
  let textContent = message.content;
  // Get the original text content
  if (typeof textContent !== "string") {
    logger.warn(
      "Message content is not a string - shouldn't happen but using message as-is",
    );
    return message;
  }

  // Process text file attachments - replace placeholder tags with full content
  for (const filePath of attachmentPaths) {
    const fileName = path.basename(filePath);
    textContent = await replaceTextAttachmentWithContent(
      textContent,
      filePath,
      fileName,
    );
  }

  // For user messages with attachments, create a content array
  const contentParts: (TextPart | ImagePart)[] = [];

  // Add the text part first with possibly modified content
  contentParts.push({
    type: "text",
    text: textContent,
  });

  // Add image parts for any image attachments
  for (const filePath of attachmentPaths) {
    const ext = path.extname(filePath).toLowerCase();
    if ([".jpg", ".jpeg", ".png", ".gif", ".webp"].includes(ext)) {
      try {
        // Read the file as a buffer and convert to base64 string
        // Using base64 strings instead of raw Buffers ensures proper JSON serialization
        // for storage in aiMessagesJson (raw Buffers serialize inefficiently and exceed size limits)
        const imageBuffer = await readFile(filePath);
        const mimeType =
          ext === ".jpg" ? "image/jpeg" : `image/${ext.slice(1)}`;
        const base64Data = imageBuffer.toString("base64");

        // Add the image to the content parts with base64 data and mediaType
        contentParts.push({
          type: "image",
          image: base64Data,
          mediaType: mimeType,
        });

        logger.log(`Added image attachment: ${filePath}`);
      } catch (error) {
        logger.error(`Error reading image file: ${error}`);
      }
    }
  }

  // Return the message with the content array
  return {
    role: "user",
    content: contentParts,
  };
}

function removeNonEssentialTags(text: string): string {
  return removeProblemReportTags(removeThinkingTags(text));
}

function removeThinkingTags(text: string): string {
  const thinkRegex = /<think>([\s\S]*?)<\/think>/g;
  return text.replace(thinkRegex, "").trim();
}

export function removeProblemReportTags(text: string): string {
  const problemReportRegex =
    /<dyad-problem-report[^>]*>[\s\S]*?<\/dyad-problem-report>/g;
  return text.replace(problemReportRegex, "").trim();
}

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;
}

function escapeDyadTags(text: string): string {
  // Escape dyad tags in reasoning content
  // We are replacing the opening tag with a look-alike character
  // to avoid issues where thinking content includes dyad tags
  // and are mishandled by:
  // 1. FE markdown parser
  // 2. Main process response processor
  return text.replace(/<dyad/g, "＜dyad").replace(/<\/dyad/g, "＜/dyad");
}

const CODEBASE_PROMPT_PREFIX = "This is my codebase.";
function createCodebasePrompt(codebaseInfo: string): string {
  return `${CODEBASE_PROMPT_PREFIX} ${codebaseInfo}`;
}

function createOtherAppsCodebasePrompt(otherAppsCodebaseInfo: string): string {
  return `
# Referenced Apps

These are the other apps that I've mentioned in my prompt. These other apps' codebases are READ-ONLY.

${otherAppsCodebaseInfo}
`;
}

async function getMcpTools(event: IpcMainInvokeEvent): Promise<ToolSet> {
  const mcpToolSet: ToolSet = {};
  try {
    const servers = await db
      .select()
      .from(mcpServers)
      .where(eq(mcpServers.enabled, true as any));
    for (const s of servers) {
      const client = await mcpManager.getClient(s.id);
      const toolSet = await client.tools();
      for (const [name, mcpTool] of Object.entries(toolSet)) {
        const key = `${String(s.name || "").replace(/[^a-zA-Z0-9_-]/g, "-")}__${String(name).replace(/[^a-zA-Z0-9_-]/g, "-")}`;
        mcpToolSet[key] = {
          description: mcpTool.description,
          inputSchema: mcpTool.inputSchema,
          execute: async (args: unknown, execCtx: ToolExecutionOptions) => {
            const inputPreview =
              typeof args === "string"
                ? args
                : Array.isArray(args)
                  ? args.join(" ")
                  : JSON.stringify(args).slice(0, 500);
            const ok = await requireMcpToolConsent(event, {
              serverId: s.id,
              serverName: s.name,
              toolName: name,
              toolDescription: mcpTool.description,
              inputPreview,
            });

            if (!ok) throw new Error(`User declined running tool ${key}`);
            const res = await mcpTool.execute(args, execCtx);

            return typeof res === "string" ? res : JSON.stringify(res);
          },
        };
      }
    }
  } catch (e) {
    logger.warn("Failed building MCP toolset", e);
  }
  return mcpToolSet;
}
