Merge branch 'main' of https://github.com/trycompai/comp into mariano/cloud-tests-bugged

Marfuen · Marfuen · commit 1c22dff5fdc0 · 2026-02-20T16:55:10.000-05:00
# Conflicts:
#	packages/docs/openapi.json
diff --git a/apps/portal/src/app/actions/login.ts b/apps/portal/src/app/actions/login.ts
@@ -55,7 +55,6 @@ export const login = createSafeActionClient({ handleServerError })
         email: parsedInput.email,
         otp: parsedInput.otp,
       },
-      asResponse: true,
     });
 
     return {
diff --git a/packages/device-agent/src/main/auth.ts b/packages/device-agent/src/main/auth.ts
@@ -120,21 +120,25 @@ export async function performLogin(deviceInfo: DeviceInfo): Promise<StoredAuth |
         authWindow.hide();
       }
 
-      // Wait for cookies to settle
-      await new Promise((r) => setTimeout(r, 500));
-
       try {
-        const authData = await extractAuthAndRegisterAll(deviceInfo);
+        const authData = await extractAuthAndRegisterAll(portalUrl, deviceInfo);
         if (authData) {
           setAuth(authData);
           log(`Auth complete: ${authData.organizations.length} org(s) registered`);
           finish(authData);
         } else {
-          log('Auth extraction returned null, closing window');
+          // extractAuthAndRegisterAll already showed an error dialog
           finish(null);
         }
       } catch (error) {
         log(`Auth extraction failed: ${error}`, 'ERROR');
+        dialog.showMessageBoxSync({
+          type: 'error',
+          title: 'Sign-In Failed',
+          message: 'An unexpected error occurred during sign-in.',
+          detail: `${error}`,
+          buttons: ['OK'],
+        });
         finish(null);
       }
     };
@@ -157,24 +161,56 @@ export async function performLogin(deviceInfo: DeviceInfo): Promise<StoredAuth |
   });
 }
 
+/**
+ * Waits for the session cookie to appear in the Electron session.
+ * Retries up to maxAttempts times with a delay between attempts.
+ */
+async function waitForSessionCookie(
+  portalUrl: string,
+  maxAttempts = 10,
+  delayMs = 500,
+): Promise<Electron.Cookie | null> {
+  for (let attempt = 1; attempt <= maxAttempts; attempt++) {
+    await new Promise((r) => setTimeout(r, delayMs));
+
+    const cookies = await session.defaultSession.cookies.get({ url: portalUrl });
+    const sessionCookie = cookies.find(
+      (c) =>
+        c.name === 'better-auth.session_token' || c.name === '__Secure-better-auth.session_token',
+    );
+
+    if (sessionCookie) {
+      log(`Session cookie found on attempt ${attempt}/${maxAttempts}`);
+      return sessionCookie;
+    }
+
+    log(`Session cookie not found yet (attempt ${attempt}/${maxAttempts})`, 'WARN');
+  }
+
+  return null;
+}
+
 /**
  * After login, fetches the user's orgs and registers the device for all of them.
- * Shows an error dialog if the user has no organizations.
+ * Shows error dialogs on failure so the user knows what went wrong.
  */
 async function extractAuthAndRegisterAll(
+  portalUrl: string,
   deviceInfo: DeviceInfo,
 ): Promise<StoredAuth | null> {
-  const portalUrl = getPortalUrl();
-
-  // 1. Get session cookie
-  const cookies = await session.defaultSession.cookies.get({ url: portalUrl });
-  const sessionCookie = cookies.find(
-    (c) =>
-      c.name === 'better-auth.session_token' || c.name === '__Secure-better-auth.session_token',
-  );
+  // 1. Wait for session cookie with retries
+  const sessionCookie = await waitForSessionCookie(portalUrl);
 
   if (!sessionCookie) {
-    log('No session cookie found after login', 'WARN');
+    log('No session cookie found after login (exhausted retries)', 'ERROR');
+    dialog.showMessageBoxSync({
+      type: 'error',
+      title: 'Sign-In Failed',
+      message: 'Could not complete sign-in.',
+      detail:
+        'The session cookie was not set after authentication. Please try again. If the problem persists, restart the app.',
+      buttons: ['OK'],
+    });
     return null;
   }
 
@@ -188,14 +224,28 @@ async function extractAuthAndRegisterAll(
 
   if (!sessionResponse.ok) {
     log(`Session fetch failed: ${sessionResponse.status}`, 'ERROR');
+    dialog.showMessageBoxSync({
+      type: 'error',
+      title: 'Sign-In Failed',
+      message: 'Could not verify your session.',
+      detail: `The server returned status ${sessionResponse.status}. Please try signing in again.`,
+      buttons: ['OK'],
+    });
     return null;
   }
 
   const sessionData = await sessionResponse.json();
   const userId = sessionData?.user?.id;
 
   if (!userId) {
-    log('No userId in session', 'ERROR');
+    log('No userId in session response', 'ERROR');
+    dialog.showMessageBoxSync({
+      type: 'error',
+      title: 'Sign-In Failed',
+      message: 'Could not retrieve your user information.',
+      detail: 'The session is missing user data. Please try signing in again.',
+      buttons: ['OK'],
+    });
     return null;
   }
 
@@ -208,6 +258,13 @@ async function extractAuthAndRegisterAll(
 
   if (!orgsResponse.ok) {
     log(`Failed to fetch organizations: ${orgsResponse.status}`, 'ERROR');
+    dialog.showMessageBoxSync({
+      type: 'error',
+      title: 'Sign-In Failed',
+      message: 'Could not fetch your organizations.',
+      detail: `The server returned status ${orgsResponse.status}. Please try signing in again.`,
+      buttons: ['OK'],
+    });
     return null;
   }
 
diff --git a/packages/device-agent/src/main/index.ts b/packages/device-agent/src/main/index.ts
@@ -7,7 +7,13 @@ import { performLogin, performLogout } from './auth';
 import { initAutoLaunch } from './auto-launch';
 import { getDeviceInfo } from './device-info';
 import { log } from './logger';
-import { runChecksNow, setSessionExpiredHandler, startScheduler, stopScheduler } from './scheduler';
+import {
+  runChecksNow,
+  setDevicesNotFoundHandler,
+  setSessionExpiredHandler,
+  startScheduler,
+  stopScheduler,
+} from './scheduler';
 import { clearAuth, getAuth, getLastCheckResults } from './store';
 import {
   createTray,
@@ -49,8 +55,11 @@ if (process.platform === 'darwin') {
 
 let currentStatus: TrayStatus = 'unauthenticated';
 let currentResults: CheckResult[] = [];
+let isSigningIn = false;
 
-// Handle session expiry: clear auth, update UI, and re-prompt login
+// Handle session expiry: clear auth, update UI, and re-prompt login.
+// stopScheduler() is called BEFORE triggerSignIn() so no periodic checks
+// can re-trigger the expired handler during the sign-in flow.
 setSessionExpiredHandler(async () => {
   log('Session expired — clearing auth and prompting re-login');
   stopScheduler();
@@ -59,23 +68,46 @@ setSessionExpiredHandler(async () => {
   currentResults = [];
   setStatus('unauthenticated');
   notifyRenderer(IPC_CHANNELS.AUTH_STATE_CHANGED, false);
-  // Auto-open sign-in so the user can re-authenticate immediately
+  triggerSignIn();
+});
+
+// Handle stale device IDs: all orgs returned 404 for the stored device IDs.
+// This happens when devices are deleted from the portal. Re-login to re-register.
+setDevicesNotFoundHandler(async () => {
+  log('All devices returned 404 — clearing auth and re-registering');
+  stopScheduler();
+  await performLogout();
+  clearAuth();
+  currentResults = [];
+  setStatus('unauthenticated');
+  notifyRenderer(IPC_CHANNELS.AUTH_STATE_CHANGED, false);
   triggerSignIn();
 });
 
 async function triggerSignIn(): Promise<void> {
+  if (isSigningIn) {
+    log('Sign-in already in progress, skipping duplicate request');
+    return;
+  }
+
+  isSigningIn = true;
   log('Sign-in flow triggered');
-  const deviceInfo = getDeviceInfo();
-  const auth = await performLogin(deviceInfo);
 
-  if (auth) {
-    const orgNames = auth.organizations.map((o) => o.organizationName).join(', ');
-    log(`Login successful: ${auth.organizations.length} org(s) — ${orgNames}`);
-    notifyRenderer(IPC_CHANNELS.AUTH_STATE_CHANGED, true);
-    setStatus('checking');
-    startScheduler(handleCheckComplete);
-  } else {
-    log('Login cancelled or failed');
+  try {
+    const deviceInfo = getDeviceInfo();
+    const auth = await performLogin(deviceInfo);
+
+    if (auth) {
+      const orgNames = auth.organizations.map((o) => o.organizationName).join(', ');
+      log(`Login successful: ${auth.organizations.length} org(s) — ${orgNames}`);
+      notifyRenderer(IPC_CHANNELS.AUTH_STATE_CHANGED, true);
+      setStatus('checking');
+      startScheduler(handleCheckComplete);
+    } else {
+      log('Login cancelled or failed');
+    }
+  } finally {
+    isSigningIn = false;
   }
 }
 
diff --git a/packages/device-agent/src/main/reporter.ts b/packages/device-agent/src/main/reporter.ts
@@ -8,6 +8,8 @@ export interface ReportResult {
   isCompliant: boolean;
   /** True if ANY org returned 401 — session has expired */
   sessionExpired: boolean;
+  /** True if ALL orgs returned 404 — stored device IDs are stale */
+  allDevicesNotFound: boolean;
 }
 
 /**
@@ -18,7 +20,7 @@ export async function reportCheckResults(checks: CheckResult[]): Promise<ReportR
   const auth = getAuth();
   if (!auth) {
     log('Cannot report check results: not authenticated', 'ERROR');
-    return { allSucceeded: false, isCompliant: false, sessionExpired: false };
+    return { allSucceeded: false, isCompliant: false, sessionExpired: false, allDevicesNotFound: false };
   }
 
   const portalUrl = getPortalUrl();
@@ -28,6 +30,7 @@ export async function reportCheckResults(checks: CheckResult[]): Promise<ReportR
   let allSucceeded = true;
   let anyNonCompliant = false;
   let sessionExpired = false;
+  let notFoundCount = 0;
 
   for (const org of auth.organizations) {
     const payload: CheckInRequest = {
@@ -59,6 +62,9 @@ export async function reportCheckResults(checks: CheckResult[]): Promise<ReportR
           sessionExpired = true;
           break; // No point trying other orgs with the same expired token
         }
+        if (response.status === 404) {
+          notFoundCount++;
+        }
         continue;
       }
 
@@ -80,5 +86,6 @@ export async function reportCheckResults(checks: CheckResult[]): Promise<ReportR
     allSucceeded,
     isCompliant: !anyNonCompliant && allSucceeded,
     sessionExpired,
+    allDevicesNotFound: notFoundCount > 0 && notFoundCount === auth.organizations.length,
   };
 }
diff --git a/packages/device-agent/src/main/scheduler.ts b/packages/device-agent/src/main/scheduler.ts
@@ -10,8 +10,10 @@ let isRunning = false;
 
 type CheckCallback = (results: CheckResult[], isCompliant: boolean) => void;
 type SessionExpiredCallback = () => void;
+type DevicesNotFoundCallback = () => void;
 
 let onSessionExpired: SessionExpiredCallback | null = null;
+let onDevicesNotFound: DevicesNotFoundCallback | null = null;
 
 /**
  * Registers a callback for when the session token is rejected (401).
@@ -20,6 +22,13 @@ export function setSessionExpiredHandler(handler: SessionExpiredCallback): void
   onSessionExpired = handler;
 }
 
+/**
+ * Registers a callback for when all device IDs return 404 (stale registrations).
+ */
+export function setDevicesNotFoundHandler(handler: DevicesNotFoundCallback): void {
+  onDevicesNotFound = handler;
+}
+
 /**
  * Starts the periodic compliance check scheduler.
  * Runs an initial check immediately, then repeats on the configured interval.
@@ -82,14 +91,20 @@ async function runChecksAndReport(onCheckComplete: CheckCallback): Promise<void>
     setLastCheckResults(results);
 
     // Report to all organizations
-    const { isCompliant, sessionExpired } = await reportCheckResults(results);
+    const { isCompliant, sessionExpired, allDevicesNotFound } = await reportCheckResults(results);
 
     if (sessionExpired) {
       log('Session expired during check-in, triggering re-authentication');
       onSessionExpired?.();
       return;
     }
 
+    if (allDevicesNotFound) {
+      log('All device IDs returned 404, triggering re-registration');
+      onDevicesNotFound?.();
+      return;
+    }
+
     log(`Check complete: ${isCompliant ? 'COMPLIANT' : 'NON-COMPLIANT'}`);
     onCheckComplete(results, isCompliant);
   } catch (error) {
diff --git a/packages/docs/openapi.json b/packages/docs/openapi.json
@@ -14390,60 +14390,6 @@
         ]
       }
     },
-    "/v1/cloud-security/trigger/{connectionId}": {
-      "post": {
-        "operationId": "CloudSecurityController_triggerScan_v1",
-        "parameters": [
-          {
-            "name": "connectionId",
-            "required": true,
-            "in": "path",
-            "schema": {
-              "type": "string"
-            }
-          }
-        ],
-        "responses": {
-          "201": {
-            "description": ""
-          }
-        },
-        "tags": [
-          "CloudSecurity"
-        ]
-      }
-    },
-    "/v1/cloud-security/runs/{runId}": {
-      "get": {
-        "operationId": "CloudSecurityController_getRunStatus_v1",
-        "parameters": [
-          {
-            "name": "runId",
-            "required": true,
-            "in": "path",
-            "schema": {
-              "type": "string"
-            }
-          },
-          {
-            "name": "connectionId",
-            "required": true,
-            "in": "query",
-            "schema": {
-              "type": "string"
-            }
-          }
-        ],
-        "responses": {
-          "200": {
-            "description": ""
-          }
-        },
-        "tags": [
-          "CloudSecurity"
-        ]
-      }
-    },
     "/v1/browserbase/org-context": {
       "post": {
         "description": "Gets the existing browser context for the org or creates a new one",
