improve login command

packages/db/package.json

@@ -53,6 +53,7 @@
   },
   "dependencies": {
     "@libsql/client": "^0.4.3",
+    "async-listen": "^3.0.1",
     "deep-diff": "^1.0.2",
     "drizzle-orm": "^0.28.6",
     "kleur": "^4.1.5",

packages/db/src/core/cli/commands/login/index.ts

@@ -1,49 +1,65 @@
-import { mkdir, writeFile } from 'node:fs/promises';
-import { createServer } from 'node:http';
 import type { AstroConfig } from 'astro';
+import { listen } from 'async-listen';
 import { cyan } from 'kleur/colors';
+import { mkdir, writeFile } from 'node:fs/promises';
+import { createServer as _createServer } from 'node:http';
 import open from 'open';
 import ora from 'ora';
 import type { Arguments } from 'yargs-parser';
 import { SESSION_LOGIN_FILE } from '../../../tokens.js';
 import { getAstroStudioUrl } from '../../../utils.js';
 
-function serveAndResolveSession(): Promise<string> {
+// NOTE(fks): How the Astro CLI login process works:
+// 1. The Astro CLI creates a temporary server to listen for the session token
+// 2. The user is directed to studio.astro.build/ to login
+// 3. The user is redirected back to the temporary server with their session token
+// 4. The temporary server receives and saves the session token, logging the user in
+// 5. The user is redirected one last time to a success/failure page
+async function createServer(): Promise<{ url: string; promise: Promise<string> }> {
 	let resolve: (value: string | PromiseLike<string>) => void,
-		reject: (value?: string | PromiseLike<string>) => void;
-	const sessionPromise = new Promise<string>((_resolve, _reject) => {
-		resolve = _resolve;
-		reject = _reject;
-	});
+		reject: (reason?: Error) => void;
 
-	const server = createServer((req, res) => {
-		res.writeHead(200);
-		res.end();
+	const server = _createServer((req, res) => {
+		// Handle the request
 		const url = new URL(req.url ?? '/', `http://${req.headers.host}`);
-		const session = url.searchParams.get('session');
-		if (!session) {
-			reject();
+		const sessionParam = url.searchParams.get('session');
+		// Handle the response & resolve the promise
+		res.statusCode = 302;
+		if (!sessionParam) {
+			res.setHeader('location', getAstroStudioUrl() + '/auth/cli/error');
+			reject(new Error('Failed to log in'));
 		} else {
-			resolve(session);
+			res.setHeader('location', getAstroStudioUrl() + '/auth/cli/success');
+			resolve(sessionParam);
 		}
-	}).listen(5710, 'localhost');
+		res.end();
+	});
 
-	return sessionPromise.finally(() => {
+	const { port } = await listen(server, 0, '127.0.0.1');
+	const serverUrl = `http://localhost:${port}`;
+	const sessionPromise = new Promise<string>((_resolve, _reject) => {
+		resolve = _resolve;
+		reject = _reject;
+	}).finally(() => {
 		server.closeAllConnections();
 		server.close();
 	});
+
+	return { url: serverUrl, promise: sessionPromise };
 }
 
 export async function cmd({ flags }: { config: AstroConfig; flags: Arguments }) {
 	let session = flags.session;
-	const loginUrl = getAstroStudioUrl() + '/auth/cli';
 
 	if (!session) {
-		console.log(`Opening ${cyan(loginUrl)} in your browser...`);
+		const { url, promise } = await createServer();
+		const loginUrl = getAstroStudioUrl() + '/auth/cli/login?returnTo=' + encodeURIComponent(url);
+		console.log(`Opening the following URL in your browser...`);
+		console.log(cyan(loginUrl));
 		console.log(`If something goes wrong, copy-and-paste the URL into your browser.`);
 		open(loginUrl);
 		const spinner = ora('Waiting for confirmation...');
-		session = await serveAndResolveSession();
+		session = await promise;
 		spinner.succeed('Successfully logged in!');
 	}
 

packages/db/src/core/cli/index.ts

@@ -8,6 +8,17 @@ export async function cli({ flags, config }: { flags: Arguments; config: AstroCo
 	// are also handled by this package, so first check if this is a db command.
 	const command = args[2] === 'db' ? args[3] : args[2];
 
+	switch (command) {
+		case 'login': {
+			const { cmd } = await import('./commands/login/index.js');
+			return await cmd({ config, flags });
+		}
+		case 'logout': {
+			const { cmd } = await import('./commands/logout/index.js');
+			return await cmd({ config, flags });
+		}
+	}
+
 	if (!config.db?.studio) {
 		console.log(STUDIO_CONFIG_MISSING_CLI_ERROR);
 		process.exit(1);
@@ -31,14 +42,6 @@ export async function cli({ flags, config }: { flags: Arguments; config: AstroCo
 			const { cmd } = await import('./commands/verify/index.js');
 			return await cmd({ config, flags });
 		}
-		case 'login': {
-			const { cmd } = await import('./commands/login/index.js');
-			return await cmd({ config, flags });
-		}
-		case 'logout': {
-			const { cmd } = await import('./commands/logout/index.js');
-			return await cmd({ config, flags });
-		}
 		case 'link': {
 			const { cmd } = await import('./commands/link/index.js');
 			return await cmd({ config, flags });