import { ipcMain, type IpcMainInvokeEvent } from "electron"; import type { CodeProposal, FileChange, ProposalResult, SqlQuery, ActionProposal, } from "../../lib/schemas"; import { db } from "../../db"; import { messages, chats } from "../../db/schema"; import { desc, eq, and, Update } from "drizzle-orm"; import path from "node:path"; // Import path for basename // Import tag parsers import { getDyadAddDependencyTags, getDyadChatSummaryTag, getDyadDeleteTags, getDyadExecuteSqlTags, getDyadRenameTags, getDyadWriteTags, processFullResponseActions, } from "../processors/response_processor"; import log from "electron-log"; import { isServerFunction } from "../../supabase_admin/supabase_utils"; import { estimateMessagesTokens, estimateTokens, getContextWindow, } from "../utils/token_utils"; import { extractCodebase } from "../../utils/codebase"; import { getDyadAppPath } from "../../paths/paths"; import { withLock } from "../utils/lock_utils"; const logger = log.scope("proposal_handlers"); // Placeholder Proposal data (can be removed or kept for reference) // const placeholderProposal: Proposal = { ... }; // Type guard for the parsed proposal structure interface ParsedProposal { title: string; files: string[]; } function isParsedProposal(obj: any): obj is ParsedProposal { return ( obj && typeof obj === "object" && typeof obj.title === "string" && Array.isArray(obj.files) && obj.files.every((file: any) => typeof file === "string") ); } // Cache for codebase token counts interface CodebaseTokenCache { chatId: number; messageId: number; messageContent: string; tokenCount: number; timestamp: number; } // Cache expiration time (5 minutes) const CACHE_EXPIRATION_MS = 5 * 60 * 1000; // In-memory cache for codebase token counts const codebaseTokenCache = new Map(); // Function to clean up expired cache entries function cleanupExpiredCacheEntries() { const now = Date.now(); let expiredCount = 0; codebaseTokenCache.forEach((entry, key) => { if (now - entry.timestamp > CACHE_EXPIRATION_MS) { codebaseTokenCache.delete(key); expiredCount++; } }); if (expiredCount > 0) { logger.log( `Cleaned up ${expiredCount} expired codebase token cache entries` ); } } // Function to get cached token count or calculate and cache it async function getCodebaseTokenCount( chatId: number, messageId: number, messageContent: string, appPath: string ): Promise { // Clean up expired cache entries first cleanupExpiredCacheEntries(); const cacheEntry = codebaseTokenCache.get(chatId); const now = Date.now(); // Check if cache is valid - same chat, message and content, and not expired if ( cacheEntry && cacheEntry.messageId === messageId && cacheEntry.messageContent === messageContent && now - cacheEntry.timestamp < CACHE_EXPIRATION_MS ) { logger.log(`Using cached codebase token count for chatId: ${chatId}`); return cacheEntry.tokenCount; } // Calculate and cache the token count logger.log(`Calculating codebase token count for chatId: ${chatId}`); const codebase = await extractCodebase(getDyadAppPath(appPath)); const tokenCount = estimateTokens(codebase); // Store in cache codebaseTokenCache.set(chatId, { chatId, messageId, messageContent, tokenCount, timestamp: now, }); return tokenCount; } const getProposalHandler = async ( _event: IpcMainInvokeEvent, { chatId }: { chatId: number } ): Promise => { return withLock("get-proposal:" + chatId, async () => { logger.log(`IPC: get-proposal called for chatId: ${chatId}`); try { // Find the latest ASSISTANT message for the chat const latestAssistantMessage = await db.query.messages.findFirst({ where: and(eq(messages.chatId, chatId), eq(messages.role, "assistant")), orderBy: [desc(messages.createdAt)], columns: { id: true, // Fetch the ID content: true, // Fetch the content to parse approvalState: true, }, }); if ( latestAssistantMessage?.content && latestAssistantMessage.id && !latestAssistantMessage?.approvalState ) { const messageId = latestAssistantMessage.id; // Get the message ID logger.log( `Found latest assistant message (ID: ${messageId}), parsing content...` ); const messageContent = latestAssistantMessage.content; const proposalTitle = getDyadChatSummaryTag(messageContent); const proposalWriteFiles = getDyadWriteTags(messageContent); const proposalRenameFiles = getDyadRenameTags(messageContent); const proposalDeleteFiles = getDyadDeleteTags(messageContent); const proposalExecuteSqlQueries = getDyadExecuteSqlTags(messageContent); const packagesAdded = getDyadAddDependencyTags(messageContent); const filesChanged = [ ...proposalWriteFiles.map((tag) => ({ name: path.basename(tag.path), path: tag.path, summary: tag.description ?? "(no change summary found)", // Generic summary type: "write" as const, isServerFunction: isServerFunction(tag.path), })), ...proposalRenameFiles.map((tag) => ({ name: path.basename(tag.to), path: tag.to, summary: `Rename from ${tag.from} to ${tag.to}`, type: "rename" as const, isServerFunction: isServerFunction(tag.to), })), ...proposalDeleteFiles.map((tag) => ({ name: path.basename(tag), path: tag, summary: `Delete file`, type: "delete" as const, isServerFunction: isServerFunction(tag), })), ]; // Check if we have enough information to create a proposal if ( filesChanged.length > 0 || packagesAdded.length > 0 || proposalExecuteSqlQueries.length > 0 ) { const proposal: CodeProposal = { type: "code-proposal", // Use parsed title or a default title if summary tag is missing but write tags exist title: proposalTitle ?? "Proposed File Changes", securityRisks: [], // Keep empty filesChanged, packagesAdded, sqlQueries: proposalExecuteSqlQueries.map((query) => ({ content: query.content, description: query.description, })), }; logger.log( "Generated code proposal. title=", proposal.title, "files=", proposal.filesChanged.length, "packages=", proposal.packagesAdded.length ); return { proposal: proposal, chatId, messageId, }; } else { logger.log( "No relevant tags found in the latest assistant message content." ); } } const actions: ActionProposal["actions"] = []; if (latestAssistantMessage?.content) { const writeTags = getDyadWriteTags(latestAssistantMessage.content); const refactorTarget = writeTags.reduce((largest, tag) => { const lineCount = tag.content.split("\n").length; return lineCount > 500 && (!largest || lineCount > largest.lineCount) ? { path: tag.path, lineCount } : largest; }, null as { path: string; lineCount: number } | null); if (refactorTarget) { actions.push({ id: "refactor-file", path: refactorTarget.path, }); } if ( writeTags.length === 0 && latestAssistantMessage.content.includes("```") ) { actions.push({ id: "write-code-properly", }); } } // Get all chat messages to calculate token usage const chat = await db.query.chats.findFirst({ where: eq(chats.id, chatId), with: { app: true, messages: { orderBy: (messages, { asc }) => [asc(messages.createdAt)], }, }, }); if (latestAssistantMessage && chat) { // Calculate total tokens from message history const messagesTokenCount = estimateMessagesTokens(chat.messages); // Use cached token count or calculate new one const codebaseTokenCount = await getCodebaseTokenCount( chatId, latestAssistantMessage.id, latestAssistantMessage.content || "", chat.app.path ); const totalTokens = messagesTokenCount + codebaseTokenCount; const contextWindow = Math.min(getContextWindow(), 100_000); logger.log( `Token usage: ${totalTokens}/${contextWindow} (${ (totalTokens / contextWindow) * 100 }%)` ); // If we're using more than 80% of the context window, suggest summarizing if (totalTokens > contextWindow * 0.8) { logger.log( `Token usage high (${totalTokens}/${contextWindow}), suggesting summarize action` ); actions.push({ id: "summarize-in-new-chat", }); } } if (actions.length > 0 && latestAssistantMessage) { return { proposal: { type: "action-proposal", actions: actions, }, chatId, messageId: latestAssistantMessage.id, }; } return null; } catch (error) { logger.error(`Error processing proposal for chatId ${chatId}:`, error); return null; // Indicate DB or processing error } }); }; // Handler to approve a proposal (process actions and update message) const approveProposalHandler = async ( _event: IpcMainInvokeEvent, { chatId, messageId }: { chatId: number; messageId: number } ): Promise<{ success: boolean; error?: string }> => { logger.log( `IPC: approve-proposal called for chatId: ${chatId}, messageId: ${messageId}` ); try { // 1. Fetch the specific assistant message const messageToApprove = await db.query.messages.findFirst({ where: and( eq(messages.id, messageId), eq(messages.chatId, chatId), eq(messages.role, "assistant") ), columns: { content: true, }, }); if (!messageToApprove?.content) { logger.error( `Assistant message not found for chatId: ${chatId}, messageId: ${messageId}` ); return { success: false, error: "Assistant message not found." }; } // 2. Process the actions defined in the message content const chatSummary = getDyadChatSummaryTag(messageToApprove.content); const processResult = await processFullResponseActions( messageToApprove.content, chatId, { chatSummary: chatSummary ?? undefined, messageId, } // Pass summary if found ); if (processResult.error) { logger.error( `Error processing actions for message ${messageId}:`, processResult.error ); // Optionally: Update message state to 'error' or similar? // For now, just return error to frontend return { success: false, error: `Action processing failed: ${processResult.error}`, }; } return { success: true }; } catch (error) { logger.error(`Error approving proposal for messageId ${messageId}:`, error); return { success: false, error: (error as Error)?.message || "Unknown error", }; } }; // Handler to reject a proposal (just update message state) const rejectProposalHandler = async ( _event: IpcMainInvokeEvent, { chatId, messageId }: { chatId: number; messageId: number } ): Promise<{ success: boolean; error?: string }> => { logger.log( `IPC: reject-proposal called for chatId: ${chatId}, messageId: ${messageId}` ); try { // 1. Verify the message exists and is an assistant message const messageToReject = await db.query.messages.findFirst({ where: and( eq(messages.id, messageId), eq(messages.chatId, chatId), eq(messages.role, "assistant") ), columns: { id: true }, }); if (!messageToReject) { logger.error( `Assistant message not found for chatId: ${chatId}, messageId: ${messageId}` ); return { success: false, error: "Assistant message not found." }; } // 2. Update the message's approval state to 'rejected' await db .update(messages) .set({ approvalState: "rejected" }) .where(eq(messages.id, messageId)); logger.log(`Message ${messageId} marked as rejected.`); return { success: true }; } catch (error) { logger.error(`Error rejecting proposal for messageId ${messageId}:`, error); return { success: false, error: (error as Error)?.message || "Unknown error", }; } }; // Function to register proposal-related handlers export function registerProposalHandlers() { ipcMain.handle("get-proposal", getProposalHandler); ipcMain.handle("approve-proposal", approveProposalHandler); ipcMain.handle("reject-proposal", rejectProposalHandler); }