Neon / portal template support (#713)

TODOs:
- [x] Do restart when checkout / restore if there is a DB
- [x] List all branches (branch id, name, date)
- [x] Allow checking out versions with no DB
- [x] safeguard to never delete main branches
- [x] create app hook for neon template
- [x] weird UX with connector on configure panel
- [x] tiny neon logo in connector
- [x] deploy to vercel
- [x] build forgot password page
- [x] what about email setup
- [x] lots of imgix errors
- [x] edit file - db snapshot
- [x] DYAD_DISABLE_DB_PUSH
- [ ] update portal doc
- [x] switch preview branch to be read-only endpoint
- [x] disable supabase sys prompt if neon is enabled
- [ ] https://payloadcms.com/docs/upload/storage-adapters
- [x] need to use main branch...

Phase 2?
- [x] generate DB migrations

src/ipc/handlers/app_env_vars_handlers.ts

@@ -10,7 +10,11 @@ import { apps } from "../../db/schema";
 import { eq } from "drizzle-orm";
 import { getDyadAppPath } from "../../paths/paths";
 import { GetAppEnvVarsParams, SetAppEnvVarsParams } from "../ipc_types";
-import { parseEnvFile, serializeEnvFile } from "../utils/app_env_var_utils";
+import {
+  ENV_FILE_NAME,
+  parseEnvFile,
+  serializeEnvFile,
+} from "../utils/app_env_var_utils";
 
 export function registerAppEnvVarsHandlers() {
   // Handler to get app environment variables
@@ -27,7 +31,7 @@ export function registerAppEnvVarsHandlers() {
         }
 
         const appPath = getDyadAppPath(app.path);
-        const envFilePath = path.join(appPath, ".env.local");
+        const envFilePath = path.join(appPath, ENV_FILE_NAME);
 
         // If .env.local doesn't exist, return empty array
         try {
@@ -63,7 +67,7 @@ export function registerAppEnvVarsHandlers() {
         }
 
         const appPath = getDyadAppPath(app.path);
-        const envFilePath = path.join(appPath, ".env.local");
+        const envFilePath = path.join(appPath, ENV_FILE_NAME);
 
         // Serialize environment variables to .env.local format
         const content = serializeEnvFile(envVars);

src/ipc/handlers/app_handlers.ts

@@ -8,6 +8,7 @@ import type {
   RenameBranchParams,
   CopyAppParams,
   EditAppFileReturnType,
+  RespondToAppInputParams,
 } from "../ipc_types";
 import fs from "node:fs";
 import path from "node:path";
@@ -47,6 +48,7 @@ import { safeSend } from "../utils/safe_sender";
 import { normalizePath } from "../../../shared/normalizePath";
 import { isServerFunction } from "@/supabase_admin/supabase_utils";
 import { getVercelTeamSlug } from "../utils/vercel_utils";
+import { storeDbTimestampAtCurrentVersion } from "../utils/neon_timestamp_utils";
 
 async function copyDir(
   source: string,
@@ -80,28 +82,32 @@ async function executeApp({
   appPath,
   appId,
   event, // Keep event for local-node case
+  isNeon,
 }: {
   appPath: string;
   appId: number;
   event: Electron.IpcMainInvokeEvent;
+  isNeon: boolean;
 }): Promise<void> {
   if (proxyWorker) {
     proxyWorker.terminate();
     proxyWorker = null;
   }
-  await executeAppLocalNode({ appPath, appId, event });
+  await executeAppLocalNode({ appPath, appId, event, isNeon });
 }
 
 async function executeAppLocalNode({
   appPath,
   appId,
   event,
+  isNeon,
 }: {
   appPath: string;
   appId: number;
   event: Electron.IpcMainInvokeEvent;
+  isNeon: boolean;
 }): Promise<void> {
-  const process = spawn(
+  const spawnedProcess = spawn(
     "(pnpm install && pnpm run dev --port 32100) || (npm install --legacy-peer-deps && npm run dev -- --port 32100)",
     [],
     {
@@ -113,11 +119,11 @@ async function executeAppLocalNode({
   );
 
   // Check if process spawned correctly
-  if (!process.pid) {
+  if (!spawnedProcess.pid) {
     // Attempt to capture any immediate errors if possible
     let errorOutput = "";
-    process.stderr?.on("data", (data) => (errorOutput += data));
-    await new Promise((resolve) => process.on("error", resolve)); // Wait for error event
+    spawnedProcess.stderr?.on("data", (data) => (errorOutput += data));
+    await new Promise((resolve) => spawnedProcess.on("error", resolve)); // Wait for error event
     throw new Error(
       `Failed to spawn process for app ${appId}. Error: ${
         errorOutput || "Unknown spawn error"
@@ -127,35 +133,69 @@ async function executeAppLocalNode({
 
   // Increment the counter and store the process reference with its ID
   const currentProcessId = processCounter.increment();
-  runningApps.set(appId, { process, processId: currentProcessId });
+  runningApps.set(appId, {
+    process: spawnedProcess,
+    processId: currentProcessId,
+  });
 
   // Log output
-  process.stdout?.on("data", async (data) => {
+  spawnedProcess.stdout?.on("data", async (data) => {
     const message = util.stripVTControlCharacters(data.toString());
-    logger.debug(`App ${appId} (PID: ${process.pid}) stdout: ${message}`);
+    logger.debug(
+      `App ${appId} (PID: ${spawnedProcess.pid}) stdout: ${message}`,
+    );
 
-    safeSend(event.sender, "app:output", {
-      type: "stdout",
-      message,
-      appId,
-    });
-    const urlMatch = message.match(/(https?:\/\/localhost:\d+\/?)/);
-    if (urlMatch) {
-      proxyWorker = await startProxy(urlMatch[1], {
-        onStarted: (proxyUrl) => {
-          safeSend(event.sender, "app:output", {
-            type: "stdout",
-            message: `[dyad-proxy-server]started=[${proxyUrl}] original=[${urlMatch[1]}]`,
-            appId,
-          });
-        },
+    // This is a hacky heuristic to pick up when drizzle is asking for user
+    // to select from one of a few choices. We automatically pick the first
+    // option because it's usually a good default choice. We guard this with
+    // isNeon because: 1) only Neon apps (for the official Dyad templates) should
+    // get this template and 2) it's safer to do this with Neon apps because
+    // their databases have point in time restore built-in.
+    if (isNeon && message.includes("created or renamed from another")) {
+      spawnedProcess.stdin.write(`\r\n`);
+      logger.info(
+        `App ${appId} (PID: ${spawnedProcess.pid}) wrote enter to stdin to automatically respond to drizzle push input`,
+      );
+    }
+
+    // Check if this is an interactive prompt requiring user input
+    const inputRequestPattern = /\s*›\s*\([yY]\/[nN]\)\s*$/;
+    const isInputRequest = inputRequestPattern.test(message);
+    if (isInputRequest) {
+      // Send special input-requested event for interactive prompts
+      safeSend(event.sender, "app:output", {
+        type: "input-requested",
+        message,
+        appId,
       });
+    } else {
+      // Normal stdout handling
+      safeSend(event.sender, "app:output", {
+        type: "stdout",
+        message,
+        appId,
+      });
+
+      const urlMatch = message.match(/(https?:\/\/localhost:\d+\/?)/);
+      if (urlMatch) {
+        proxyWorker = await startProxy(urlMatch[1], {
+          onStarted: (proxyUrl) => {
+            safeSend(event.sender, "app:output", {
+              type: "stdout",
+              message: `[dyad-proxy-server]started=[${proxyUrl}] original=[${urlMatch[1]}]`,
+              appId,
+            });
+          },
+        });
+      }
     }
   });
 
-  process.stderr?.on("data", (data) => {
+  spawnedProcess.stderr?.on("data", (data) => {
     const message = util.stripVTControlCharacters(data.toString());
-    logger.error(`App ${appId} (PID: ${process.pid}) stderr: ${message}`);
+    logger.error(
+      `App ${appId} (PID: ${spawnedProcess.pid}) stderr: ${message}`,
+    );
     safeSend(event.sender, "app:output", {
       type: "stderr",
       message,
@@ -164,19 +204,19 @@ async function executeAppLocalNode({
   });
 
   // Handle process exit/close
-  process.on("close", (code, signal) => {
+  spawnedProcess.on("close", (code, signal) => {
     logger.log(
-      `App ${appId} (PID: ${process.pid}) process closed with code ${code}, signal ${signal}.`,
+      `App ${appId} (PID: ${spawnedProcess.pid}) process closed with code ${code}, signal ${signal}.`,
     );
-    removeAppIfCurrentProcess(appId, process);
+    removeAppIfCurrentProcess(appId, spawnedProcess);
   });
 
   // Handle errors during process lifecycle (e.g., command not found)
-  process.on("error", (err) => {
+  spawnedProcess.on("error", (err) => {
     logger.error(
-      `Error in app ${appId} (PID: ${process.pid}) process: ${err.message}`,
+      `Error in app ${appId} (PID: ${spawnedProcess.pid}) process: ${err.message}`,
     );
-    removeAppIfCurrentProcess(appId, process);
+    removeAppIfCurrentProcess(appId, spawnedProcess);
     // Note: We don't throw here as the error is asynchronous. The caller got a success response already.
     // Consider adding ipcRenderer event emission to notify UI of the error.
   });
@@ -466,7 +506,12 @@ export function registerAppHandlers() {
         try {
           // Kill any orphaned process on port 32100 (in case previous run left it)
           await killProcessOnPort(32100);
-          await executeApp({ appPath, appId, event });
+          await executeApp({
+            appPath,
+            appId,
+            event,
+            isNeon: !!app.neonProjectId,
+          });
 
           return;
         } catch (error: any) {
@@ -596,7 +641,12 @@ export function registerAppHandlers() {
             `Executing app ${appId} in path ${app.path} after restart request`,
           ); // Adjusted log
 
-          await executeApp({ appPath, appId, event }); // This will handle starting either mode
+          await executeApp({
+            appPath,
+            appId,
+            event,
+            isNeon: !!app.neonProjectId,
+          }); // This will handle starting either mode
 
           return;
         } catch (error) {
@@ -633,6 +683,23 @@ export function registerAppHandlers() {
         throw new Error("Invalid file path");
       }
 
+      if (app.neonProjectId && app.neonDevelopmentBranchId) {
+        try {
+          await storeDbTimestampAtCurrentVersion({
+            appId: app.id,
+          });
+        } catch (error) {
+          logger.error(
+            "Error storing Neon timestamp at current version:",
+            error,
+          );
+          throw new Error(
+            "Could not store Neon timestamp at current version; database versioning functionality is not working: " +
+              error,
+          );
+        }
+      }
+
       // Ensure directory exists
       const dirPath = path.dirname(fullPath);
       await fsPromises.mkdir(dirPath, { recursive: true });
@@ -968,4 +1035,33 @@ export function registerAppHandlers() {
       }
     });
   });
+
+  handle(
+    "respond-to-app-input",
+    async (_, { appId, response }: RespondToAppInputParams) => {
+      if (response !== "y" && response !== "n") {
+        throw new Error(`Invalid response: ${response}`);
+      }
+      const appInfo = runningApps.get(appId);
+
+      if (!appInfo) {
+        throw new Error(`App ${appId} is not running`);
+      }
+
+      const { process } = appInfo;
+
+      if (!process.stdin) {
+        throw new Error(`App ${appId} process has no stdin available`);
+      }
+
+      try {
+        // Write the response to stdin with a newline
+        process.stdin.write(`${response}\n`);
+        logger.debug(`Sent response '${response}' to app ${appId} stdin`);
+      } catch (error: any) {
+        logger.error(`Error sending response to app ${appId}:`, error);
+        throw new Error(`Failed to send response to app: ${error.message}`);
+      }
+    },
+  );
 }

src/ipc/handlers/chat_stream_handlers.ts

@@ -456,7 +456,10 @@ ${componentSnippet}
             (await getSupabaseContext({
               supabaseProjectId: updatedChat.app.supabaseProjectId,
             }));
-        } else {
+        } else if (
+          // Neon projects don't need Supabase.
+          !updatedChat.app?.neonProjectId
+        ) {
           systemPrompt += "\n\n" + SUPABASE_NOT_AVAILABLE_SYSTEM_PROMPT;
         }
         const isSummarizeIntent = req.prompt.startsWith(

src/ipc/handlers/neon_handlers.ts

@@ -0,0 +1,235 @@
+import log from "electron-log";
+
+import { createTestOnlyLoggedHandler } from "./safe_handle";
+import { handleNeonOAuthReturn } from "../../neon_admin/neon_return_handler";
+import {
+  getNeonClient,
+  getNeonErrorMessage,
+  getNeonOrganizationId,
+} from "../../neon_admin/neon_management_client";
+import {
+  CreateNeonProjectParams,
+  NeonProject,
+  GetNeonProjectParams,
+  GetNeonProjectResponse,
+  NeonBranch,
+} from "../ipc_types";
+import { db } from "../../db";
+import { apps } from "../../db/schema";
+import { eq } from "drizzle-orm";
+import { ipcMain } from "electron";
+import { EndpointType } from "@neondatabase/api-client";
+import { retryOnLocked } from "../utils/retryOnLocked";
+
+export const logger = log.scope("neon_handlers");
+
+const testOnlyHandle = createTestOnlyLoggedHandler(logger);
+
+export function registerNeonHandlers() {
+  // Do not use log handler because there's sensitive data in the response
+  ipcMain.handle(
+    "neon:create-project",
+    async (
+      _,
+      { name, appId }: CreateNeonProjectParams,
+    ): Promise<NeonProject> => {
+      const neonClient = await getNeonClient();
+
+      logger.info(`Creating Neon project: ${name} for app ${appId}`);
+
+      try {
+        // Get the organization ID
+        const orgId = await getNeonOrganizationId();
+
+        // Create project with retry on locked errors
+        const response = await retryOnLocked(
+          () =>
+            neonClient.createProject({
+              project: {
+                name: name,
+                org_id: orgId,
+              },
+            }),
+          `Create project ${name} for app ${appId}`,
+        );
+
+        if (!response.data.project) {
+          throw new Error(
+            "Failed to create project: No project data returned.",
+          );
+        }
+
+        const project = response.data.project;
+        const developmentBranch = response.data.branch;
+
+        const previewBranchResponse = await retryOnLocked(
+          () =>
+            neonClient.createProjectBranch(project.id, {
+              endpoints: [{ type: EndpointType.ReadOnly }],
+              branch: {
+                name: "preview",
+                parent_id: developmentBranch.id,
+              },
+            }),
+          `Create preview branch for project ${project.id}`,
+        );
+
+        if (
+          !previewBranchResponse.data.branch ||
+          !previewBranchResponse.data.connection_uris
+        ) {
+          throw new Error(
+            "Failed to create preview branch: No branch data returned.",
+          );
+        }
+
+        const previewBranch = previewBranchResponse.data.branch;
+
+        // Store project and branch info in the app's DB row
+        await db
+          .update(apps)
+          .set({
+            neonProjectId: project.id,
+            neonDevelopmentBranchId: developmentBranch.id,
+            neonPreviewBranchId: previewBranch.id,
+          })
+          .where(eq(apps.id, appId));
+
+        logger.info(
+          `Successfully created Neon project: ${project.id} and development branch: ${developmentBranch.id} for app ${appId}`,
+        );
+        return {
+          id: project.id,
+          name: project.name,
+          connectionString: response.data.connection_uris[0].connection_uri,
+          branchId: developmentBranch.id,
+        };
+      } catch (error: any) {
+        const errorMessage = getNeonErrorMessage(error);
+        const message = `Failed to create Neon project for app ${appId}: ${errorMessage}`;
+        logger.error(message);
+        throw new Error(message);
+      }
+    },
+  );
+
+  ipcMain.handle(
+    "neon:get-project",
+    async (
+      _,
+      { appId }: GetNeonProjectParams,
+    ): Promise<GetNeonProjectResponse> => {
+      logger.info(`Getting Neon project info for app ${appId}`);
+
+      try {
+        // Get the app from the database to find the neonProjectId and neonBranchId
+        const app = await db
+          .select()
+          .from(apps)
+          .where(eq(apps.id, appId))
+          .limit(1);
+
+        if (app.length === 0) {
+          throw new Error(`App with ID ${appId} not found`);
+        }
+
+        const appData = app[0];
+        if (!appData.neonProjectId) {
+          throw new Error(`No Neon project found for app ${appId}`);
+        }
+
+        const neonClient = await getNeonClient();
+        console.log("PROJECT ID", appData.neonProjectId);
+
+        // Get project info
+        const projectResponse = await neonClient.getProject(
+          appData.neonProjectId,
+        );
+
+        if (!projectResponse.data.project) {
+          throw new Error("Failed to get project: No project data returned.");
+        }
+
+        const project = projectResponse.data.project;
+
+        // Get list of branches
+        const branchesResponse = await neonClient.listProjectBranches({
+          projectId: appData.neonProjectId,
+        });
+
+        if (!branchesResponse.data.branches) {
+          throw new Error("Failed to get branches: No branch data returned.");
+        }
+
+        // Map branches to our format
+        const branches: NeonBranch[] = branchesResponse.data.branches.map(
+          (branch) => {
+            let type: "production" | "development" | "snapshot" | "preview";
+
+            if (branch.default) {
+              type = "production";
+            } else if (branch.id === appData.neonDevelopmentBranchId) {
+              type = "development";
+            } else if (branch.id === appData.neonPreviewBranchId) {
+              type = "preview";
+            } else {
+              type = "snapshot";
+            }
+
+            // Find parent branch name if parent_id exists
+            let parentBranchName: string | undefined;
+            if (branch.parent_id) {
+              const parentBranch = branchesResponse.data.branches?.find(
+                (b) => b.id === branch.parent_id,
+              );
+              parentBranchName = parentBranch?.name;
+            }
+
+            return {
+              type,
+              branchId: branch.id,
+              branchName: branch.name,
+              lastUpdated: branch.updated_at,
+              parentBranchId: branch.parent_id,
+              parentBranchName,
+            };
+          },
+        );
+
+        logger.info(
+          `Successfully retrieved Neon project info for app ${appId}`,
+        );
+
+        return {
+          projectId: project.id,
+          projectName: project.name,
+          orgId: project.org_id ?? "<unknown_org_id>",
+          branches,
+        };
+      } catch (error) {
+        logger.error(
+          `Failed to get Neon project info for app ${appId}:`,
+          error,
+        );
+        throw error;
+      }
+    },
+  );
+
+  testOnlyHandle("neon:fake-connect", async (event) => {
+    // Call handleNeonOAuthReturn with fake data
+    handleNeonOAuthReturn({
+      token: "fake-neon-access-token",
+      refreshToken: "fake-neon-refresh-token",
+      expiresIn: 3600, // 1 hour
+    });
+    logger.info("Called handleNeonOAuthReturn with fake data during testing.");
+
+    // Simulate the deep link event
+    event.sender.send("deep-link-received", {
+      type: "neon-oauth-return",
+      url: "https://oauth.dyad.sh/api/integrations/neon/login",
+    });
+    logger.info("Sent fake neon deep-link-received event during testing.");
+  });
+}

src/ipc/handlers/portal_handlers.ts

@@ -0,0 +1,138 @@
+import { createLoggedHandler } from "./safe_handle";
+import log from "electron-log";
+import { db } from "../../db";
+import { apps } from "../../db/schema";
+import { eq } from "drizzle-orm";
+import { getDyadAppPath } from "../../paths/paths";
+import { spawn } from "child_process";
+import fs from "node:fs";
+import git from "isomorphic-git";
+import { gitCommit } from "../utils/git_utils";
+import { storeDbTimestampAtCurrentVersion } from "../utils/neon_timestamp_utils";
+
+const logger = log.scope("portal_handlers");
+const handle = createLoggedHandler(logger);
+
+async function getApp(appId: number) {
+  const app = await db.query.apps.findFirst({
+    where: eq(apps.id, appId),
+  });
+  if (!app) {
+    throw new Error(`App with id ${appId} not found`);
+  }
+  return app;
+}
+
+export function registerPortalHandlers() {
+  handle(
+    "portal:migrate-create",
+    async (_, { appId }: { appId: number }): Promise<{ output: string }> => {
+      const app = await getApp(appId);
+      const appPath = getDyadAppPath(app.path);
+
+      // Run the migration command
+      const migrationOutput = await new Promise<string>((resolve, reject) => {
+        logger.info(`Running migrate:create for app ${appId} at ${appPath}`);
+
+        const process = spawn("npm run migrate:create -- --skip-empty", {
+          cwd: appPath,
+          shell: true,
+          stdio: "pipe",
+        });
+
+        let stdout = "";
+        let stderr = "";
+
+        process.stdout?.on("data", (data) => {
+          const output = data.toString();
+          stdout += output;
+          logger.info(`migrate:create stdout: ${output}`);
+          if (output.includes("created or renamed from another")) {
+            process.stdin.write(`\r\n`);
+            logger.info(
+              `App ${appId} (PID: ${process.pid}) wrote enter to stdin to automatically respond to drizzle migrate input`,
+            );
+          }
+        });
+
+        process.stderr?.on("data", (data) => {
+          const output = data.toString();
+          stderr += output;
+          logger.warn(`migrate:create stderr: ${output}`);
+        });
+
+        process.on("close", (code) => {
+          const combinedOutput =
+            stdout + (stderr ? `\n\nErrors/Warnings:\n${stderr}` : "");
+
+          if (code === 0) {
+            if (stdout.includes("Migration created at")) {
+              logger.info(
+                `migrate:create completed successfully for app ${appId}`,
+              );
+              resolve(combinedOutput);
+            } else {
+              logger.error(
+                `migrate:create completed successfully for app ${appId} but no migration was created`,
+              );
+              reject(
+                new Error(
+                  "No migration was created because no changes were found.",
+                ),
+              );
+            }
+          } else {
+            logger.error(
+              `migrate:create failed for app ${appId} with exit code ${code}`,
+            );
+            const errorMessage = `Migration creation failed (exit code ${code})\n\n${combinedOutput}`;
+            reject(new Error(errorMessage));
+          }
+        });
+
+        process.on("error", (err) => {
+          logger.error(`Failed to spawn migrate:create for app ${appId}:`, err);
+          const errorMessage = `Failed to run migration command: ${err.message}\n\nOutput:\n${stdout}\n\nErrors:\n${stderr}`;
+          reject(new Error(errorMessage));
+        });
+      });
+
+      if (app.neonProjectId && app.neonDevelopmentBranchId) {
+        try {
+          await storeDbTimestampAtCurrentVersion({
+            appId: app.id,
+          });
+        } catch (error) {
+          logger.error(
+            "Error storing Neon timestamp at current version:",
+            error,
+          );
+          throw new Error(
+            "Could not store Neon timestamp at current version; database versioning functionality is not working: " +
+              error,
+          );
+        }
+      }
+
+      // Stage all changes and commit
+      try {
+        await git.add({
+          fs,
+          dir: appPath,
+          filepath: ".",
+        });
+
+        const commitHash = await gitCommit({
+          path: appPath,
+          message: "[dyad] Generate database migration file",
+        });
+
+        logger.info(`Successfully committed migration changes: ${commitHash}`);
+        return { output: migrationOutput };
+      } catch (gitError) {
+        logger.error(`Migration created but failed to commit: ${gitError}`);
+        throw new Error(`Migration created but failed to commit: ${gitError}`);
+      }
+    },
+  );
+}

src/ipc/handlers/token_count_handlers.ts

@@ -65,7 +65,10 @@ export function registerTokenCountHandlers() {
         supabaseContext = await getSupabaseContext({
           supabaseProjectId: chat.app.supabaseProjectId,
         });
-      } else {
+      } else if (
+        // Neon projects don't need Supabase.
+        !chat.app?.neonProjectId
+      ) {
         systemPrompt += "\n\n" + SUPABASE_NOT_AVAILABLE_SYSTEM_PROMPT;
       }
 

src/ipc/handlers/version_handlers.ts

@@ -1,7 +1,12 @@
 import { db } from "../../db";
-import { apps, messages } from "../../db/schema";
+import { apps, messages, versions } from "../../db/schema";
 import { desc, eq, and, gt } from "drizzle-orm";
-import type { Version, BranchResult } from "../ipc_types";
+import type {
+  Version,
+  BranchResult,
+  RevertVersionParams,
+  RevertVersionResponse,
+} from "../ipc_types";
 import fs from "node:fs";
 import path from "node:path";
 import { getDyadAppPath } from "../../paths/paths";
@@ -11,10 +16,51 @@ import log from "electron-log";
 import { createLoggedHandler } from "./safe_handle";
 import { gitCheckout, gitCommit, gitStageToRevert } from "../utils/git_utils";
 
+import {
+  getNeonClient,
+  getNeonErrorMessage,
+} from "../../neon_admin/neon_management_client";
+import {
+  updatePostgresUrlEnvVar,
+  updateDbPushEnvVar,
+} from "../utils/app_env_var_utils";
+import { storeDbTimestampAtCurrentVersion } from "../utils/neon_timestamp_utils";
+import { retryOnLocked } from "../utils/retryOnLocked";
+
 const logger = log.scope("version_handlers");
 
 const handle = createLoggedHandler(logger);
 
+async function restoreBranchForPreview({
+  appId,
+  dbTimestamp,
+  neonProjectId,
+  previewBranchId,
+  developmentBranchId,
+}: {
+  appId: number;
+  dbTimestamp: string;
+  neonProjectId: string;
+  previewBranchId: string;
+  developmentBranchId: string;
+}): Promise<void> {
+  try {
+    const neonClient = await getNeonClient();
+    await retryOnLocked(
+      () =>
+        neonClient.restoreProjectBranch(neonProjectId, previewBranchId, {
+          source_branch_id: developmentBranchId,
+          source_timestamp: dbTimestamp,
+        }),
+      `Restore preview branch ${previewBranchId} for app ${appId}`,
+    );
+  } catch (error) {
+    const errorMessage = getNeonErrorMessage(error);
+    logger.error("Error in restoreBranchForPreview:", errorMessage);
+    throw new Error(errorMessage);
+  }
+}
+
 export function registerVersionHandlers() {
   handle("list-versions", async (_, { appId }: { appId: number }) => {
     const app = await db.query.apps.findFirst({
@@ -40,11 +86,32 @@ export function registerVersionHandlers() {
       depth: 100_000, // Limit to last 100_000 commits for performance
     });
 
-    return commits.map((commit: ReadCommitResult) => ({
-      oid: commit.oid,
-      message: commit.commit.message,
-      timestamp: commit.commit.author.timestamp,
-    })) satisfies Version[];
+    // Get all snapshots for this app to match with commits
+    const appSnapshots = await db.query.versions.findMany({
+      where: eq(versions.appId, appId),
+    });
+
+    // Create a map of commitHash -> snapshot info for quick lookup
+    const snapshotMap = new Map<
+      string,
+      { neonDbTimestamp: string | null; createdAt: Date }
+    >();
+    for (const snapshot of appSnapshots) {
+      snapshotMap.set(snapshot.commitHash, {
+        neonDbTimestamp: snapshot.neonDbTimestamp,
+        createdAt: snapshot.createdAt,
+      });
+    }
+
+    return commits.map((commit: ReadCommitResult) => {
+      const snapshotInfo = snapshotMap.get(commit.oid);
+      return {
+        oid: commit.oid,
+        message: commit.commit.message,
+        timestamp: commit.commit.author.timestamp,
+        dbTimestamp: snapshotInfo?.neonDbTimestamp,
+      };
+    }) satisfies Version[];
   });
 
   handle(
@@ -86,12 +153,11 @@ export function registerVersionHandlers() {
     "revert-version",
     async (
       _,
-      {
-        appId,
-        previousVersionId,
-      }: { appId: number; previousVersionId: string },
-    ): Promise<void> => {
+      { appId, previousVersionId }: RevertVersionParams,
+    ): Promise<RevertVersionResponse> => {
       return withLock(appId, async () => {
+        let successMessage = "Restored version";
+        let warningMessage: string | undefined = undefined;
         const app = await db.query.apps.findFirst({
           where: eq(apps.id, appId),
         });
@@ -101,12 +167,26 @@ export function registerVersionHandlers() {
         }
 
         const appPath = getDyadAppPath(app.path);
+        // Get the current commit hash before reverting
+        const currentCommitHash = await git.resolveRef({
+          fs,
+          dir: appPath,
+          ref: "main",
+        });
 
         await gitCheckout({
           path: appPath,
           ref: "main",
         });
 
+        if (app.neonProjectId && app.neonDevelopmentBranchId) {
+          // We are going to add a new commit on top, so let's store
+          // the current timestamp at the current version.
+          await storeDbTimestampAtCurrentVersion({
+            appId,
+          });
+        }
+
         await gitStageToRevert({
           path: appPath,
           targetOid: previousVersionId,
@@ -154,6 +234,87 @@ export function registerVersionHandlers() {
               );
           }
         }
+
+        if (app.neonProjectId && app.neonDevelopmentBranchId) {
+          const version = await db.query.versions.findFirst({
+            where: and(
+              eq(versions.appId, appId),
+              eq(versions.commitHash, previousVersionId),
+            ),
+          });
+          if (version && version.neonDbTimestamp) {
+            try {
+              const preserveBranchName = `preserve_${currentCommitHash}-${Date.now()}`;
+              const neonClient = await getNeonClient();
+              const response = await retryOnLocked(
+                () =>
+                  neonClient.restoreProjectBranch(
+                    app.neonProjectId!,
+                    app.neonDevelopmentBranchId!,
+                    {
+                      source_branch_id: app.neonDevelopmentBranchId!,
+                      source_timestamp: version.neonDbTimestamp!,
+                      preserve_under_name: preserveBranchName,
+                    },
+                  ),
+                `Restore development branch ${app.neonDevelopmentBranchId} for app ${appId}`,
+              );
+              // Update all versions which have a newer DB timestamp than the version we're restoring to
+              // and remove their DB timestamp.
+              await db
+                .update(versions)
+                .set({ neonDbTimestamp: null })
+                .where(
+                  and(
+                    eq(versions.appId, appId),
+                    gt(versions.neonDbTimestamp, version.neonDbTimestamp),
+                  ),
+                );
+
+              const preserveBranchId = response.data.branch.parent_id;
+              if (!preserveBranchId) {
+                throw new Error("Preserve branch ID not found");
+              }
+              logger.info(
+                `Deleting preserve branch ${preserveBranchId} for app ${appId}`,
+              );
+              try {
+                // Intentionally do not await this because it's not
+                // critical for the restore operation, it's to clean up branches
+                // so the user doesn't hit the branch limit later.
+                retryOnLocked(
+                  () =>
+                    neonClient.deleteProjectBranch(
+                      app.neonProjectId!,
+                      preserveBranchId,
+                    ),
+                  `Delete preserve branch ${preserveBranchId} for app ${appId}`,
+                  { retryBranchWithChildError: true },
+                );
+              } catch (error) {
+                const errorMessage = getNeonErrorMessage(error);
+                logger.error("Error in deleteProjectBranch:", errorMessage);
+              }
+            } catch (error) {
+              const errorMessage = getNeonErrorMessage(error);
+              logger.error("Error in restoreBranchForCheckout:", errorMessage);
+              warningMessage = `Could not restore database because of error: ${errorMessage}`;
+              // Do not throw, so we can finish switching the postgres branch
+              // It might throw because they picked a timestamp that's too old.
+            }
+            successMessage =
+              "Successfully restored to version (including database)";
+          }
+          await switchPostgresToDevelopmentBranch({
+            neonProjectId: app.neonProjectId,
+            neonDevelopmentBranchId: app.neonDevelopmentBranchId,
+            appPath: app.path,
+          });
+        }
+        if (warningMessage) {
+          return { warningMessage };
+        }
+        return { successMessage };
       });
     },
   );
@@ -162,7 +323,7 @@ export function registerVersionHandlers() {
     "checkout-version",
     async (
       _,
-      { appId, versionId }: { appId: number; versionId: string },
+      { appId, versionId: gitRef }: { appId: number; versionId: string },
     ): Promise<void> => {
       return withLock(appId, async () => {
         const app = await db.query.apps.findFirst({
@@ -173,13 +334,106 @@ export function registerVersionHandlers() {
           throw new Error("App not found");
         }
 
-        const appPath = getDyadAppPath(app.path);
+        if (
+          app.neonProjectId &&
+          app.neonDevelopmentBranchId &&
+          app.neonPreviewBranchId
+        ) {
+          if (gitRef === "main") {
+            logger.info(
+              `Switching Postgres to development branch for app ${appId}`,
+            );
+            await switchPostgresToDevelopmentBranch({
+              neonProjectId: app.neonProjectId,
+              neonDevelopmentBranchId: app.neonDevelopmentBranchId,
+              appPath: app.path,
+            });
+          } else {
+            logger.info(
+              `Switching Postgres to preview branch for app ${appId}`,
+            );
 
+            // Regardless of whether we have a timestamp or not, we want to disable DB push
+            // while we're checking out an earlier version
+            await updateDbPushEnvVar({
+              appPath: app.path,
+              disabled: true,
+            });
+
+            const version = await db.query.versions.findFirst({
+              where: and(
+                eq(versions.appId, appId),
+                eq(versions.commitHash, gitRef),
+              ),
+            });
+
+            if (version && version.neonDbTimestamp) {
+              // SWITCH the env var for POSTGRES_URL to the preview branch
+              const neonClient = await getNeonClient();
+              const connectionUri = await neonClient.getConnectionUri({
+                projectId: app.neonProjectId,
+                branch_id: app.neonPreviewBranchId,
+                // This is the default database name for Neon
+                database_name: "neondb",
+                // This is the default role name for Neon
+                role_name: "neondb_owner",
+              });
+
+              await restoreBranchForPreview({
+                appId,
+                dbTimestamp: version.neonDbTimestamp,
+                neonProjectId: app.neonProjectId,
+                previewBranchId: app.neonPreviewBranchId,
+                developmentBranchId: app.neonDevelopmentBranchId,
+              });
+
+              await updatePostgresUrlEnvVar({
+                appPath: app.path,
+                connectionUri: connectionUri.data.uri,
+              });
+              logger.info(
+                `Switched Postgres to preview branch for app ${appId} commit ${version.commitHash} dbTimestamp=${version.neonDbTimestamp}`,
+              );
+            }
+          }
+        }
+        const fullAppPath = getDyadAppPath(app.path);
         await gitCheckout({
-          path: appPath,
-          ref: versionId,
+          path: fullAppPath,
+          ref: gitRef,
         });
       });
     },
   );
 }
+
+async function switchPostgresToDevelopmentBranch({
+  neonProjectId,
+  neonDevelopmentBranchId,
+  appPath,
+}: {
+  neonProjectId: string;
+  neonDevelopmentBranchId: string;
+  appPath: string;
+}) {
+  // SWITCH the env var for POSTGRES_URL to the development branch
+  const neonClient = await getNeonClient();
+  const connectionUri = await neonClient.getConnectionUri({
+    projectId: neonProjectId,
+    branch_id: neonDevelopmentBranchId,
+    // This is the default database name for Neon
+    database_name: "neondb",
+    // This is the default role name for Neon
+    role_name: "neondb_owner",
+  });
+
+  await updatePostgresUrlEnvVar({
+    appPath,
+    connectionUri: connectionUri.data.uri,
+  });
+
+  await updateDbPushEnvVar({
+    appPath,
+    disabled: false,
+  });
+}