fix: restore account switching on Antigravity 1.16.5 and migrate db sync to drizzle (#69)

src/ipc/database/cloudHandler.ts

@@ -2,12 +2,50 @@ import Database from 'better-sqlite3';
 import path from 'path';
 import fs from 'fs';
 import { v4 as uuidv4 } from 'uuid';
+import { desc, eq } from 'drizzle-orm';
+import type { BetterSQLite3Database } from 'drizzle-orm/better-sqlite3';
 import { getCloudAccountsDbPath, getAntigravityDbPaths } from '../../utils/paths';
 import { logger } from '../../utils/logger';
 import { CloudAccount } from '../../types/cloudAccount';
+import { ItemTableValueRowSchema, TableInfoRowSchema } from '../../types/db';
 import { decryptWithMigration, encrypt, type KeySource } from '../../utils/security';
 import { ProtobufUtils } from '../../utils/protobuf';
 import { GoogleAPIService } from '../../services/GoogleAPIService';
+import { getAntigravityVersion, isNewVersion } from '../../utils/antigravityVersion';
+import { parseRow, parseRows } from '../../utils/sqlite';
+import { configureDatabase, openDrizzleConnection } from './dbConnection';
+import { accounts, itemTable, settings } from './schema';
+import * as drizzleSchema from './schema';
+
+const SQLITE_BUSY_CODES = new Set(['SQLITE_BUSY', 'SQLITE_LOCKED']);
+const SQLITE_BUSY_TIMEOUT_MS = 3000;
+const SQLITE_RETRY_DELAY_MS = 150;
+const SQLITE_MAX_RETRIES = 3;
+
+type DrizzleExecutor = Pick<
+  BetterSQLite3Database<typeof drizzleSchema>,
+  'insert' | 'update' | 'delete' | 'select'
+>;
+
+function isSqliteBusyError(error: unknown): boolean {
+  if (!error || typeof error !== 'object') {
+    return false;
+  }
+  const err = error as { code?: string; message?: string };
+  if (err.code && SQLITE_BUSY_CODES.has(err.code)) {
+    return true;
+  }
+  if (typeof err.message === 'string') {
+    return err.message.includes('SQLITE_BUSY') || err.message.includes('SQLITE_LOCKED');
+  }
+  return false;
+}
+
+function sleepSync(ms: number): void {
+  const buffer = new SharedArrayBuffer(4);
+  const array = new Int32Array(buffer);
+  Atomics.wait(array, 0, 0, ms);
+}
 
 /**
  * Ensures that the cloud database file and schema exist.
@@ -22,7 +60,7 @@ function ensureDatabaseInitialized(dbPath: string): void {
   let db: Database.Database | null = null;
   try {
     db = new Database(dbPath);
-    db.pragma('journal_mode = WAL');
+    configureDatabase(db, { busyTimeoutMs: SQLITE_BUSY_TIMEOUT_MS });
 
     // Create accounts table
     // Storing complex objects (token, quota) as JSON strings for simplicity
@@ -43,7 +81,8 @@ function ensureDatabaseInitialized(dbPath: string): void {
     `);
 
     // Migration: Check if is_active column exists
-    const tableInfo = db.pragma('table_info(accounts)') as any[];
+    const tableInfoRaw = db.pragma('table_info(accounts)') as any[];
+    const tableInfo = parseRows(TableInfoRowSchema, tableInfoRaw, 'cloud.accounts.tableInfo');
     const hasIsActive = tableInfo.some((col) => col.name === 'is_active');
     if (!hasIsActive) {
       db.exec('ALTER TABLE accounts ADD COLUMN is_active INTEGER DEFAULT 0');
@@ -71,12 +110,28 @@ function ensureDatabaseInitialized(dbPath: string): void {
 /**
  * Gets a connection to the cloud accounts database.
  */
-function getDb(): Database.Database {
+function getCloudDb(): {
+  raw: Database.Database;
+  orm: BetterSQLite3Database<typeof drizzleSchema>;
+} {
   const dbPath = getCloudAccountsDbPath();
   ensureDatabaseInitialized(dbPath);
-  const db = new Database(dbPath);
-  db.pragma('journal_mode = WAL');
-  return db;
+  return openDrizzleConnection(
+    dbPath,
+    { readonly: false, fileMustExist: false },
+    { busyTimeoutMs: SQLITE_BUSY_TIMEOUT_MS },
+  );
+}
+
+function getIdeDb(
+  dbPath: string,
+  readOnly: boolean,
+): { raw: Database.Database; orm: BetterSQLite3Database<typeof drizzleSchema> } {
+  return openDrizzleConnection(
+    dbPath,
+    { readonly: readOnly },
+    { readOnly, busyTimeoutMs: SQLITE_BUSY_TIMEOUT_MS },
+  );
 }
 
 interface MigrationStats {
@@ -102,9 +157,9 @@ function createMigrationStats(): MigrationStats {
 }
 
 async function decryptAndMigrateField(
-  db: Database.Database,
+  orm: DrizzleExecutor,
   accountId: string,
-  field: 'token_json' | 'quota_json',
+  field: 'tokenJson' | 'quotaJson',
   value: string | null,
 ): Promise<{ value: string | null; migrated: boolean; usedFallback?: KeySource }> {
   if (!value) {
@@ -113,16 +168,18 @@ async function decryptAndMigrateField(
 
   const result = await decryptWithMigration(value);
   if (result.reencrypted) {
-    if (field === 'token_json') {
-      db.prepare('UPDATE accounts SET token_json = ? WHERE id = ?').run(
-        result.reencrypted,
-        accountId,
-      );
+    if (field === 'tokenJson') {
+      orm
+        .update(accounts)
+        .set({ tokenJson: result.reencrypted })
+        .where(eq(accounts.id, accountId))
+        .run();
     } else {
-      db.prepare('UPDATE accounts SET quota_json = ? WHERE id = ?').run(
-        result.reencrypted,
-        accountId,
-      );
+      orm
+        .update(accounts)
+        .set({ quotaJson: result.reencrypted })
+        .where(eq(accounts.id, accountId))
+        .run();
     }
     logger.info(
       `Migrated ${field} for account ${accountId} from ${result.usedFallback ?? 'unknown'} key`,
@@ -139,6 +196,8 @@ async function decryptAndMigrateField(
 type DecryptFieldResult = Awaited<ReturnType<typeof decryptAndMigrateField>>;
 
 export class CloudAccountRepo {
+  private static versionFailureLogged = false;
+
   static async init(): Promise<void> {
     const dbPath = getCloudAccountsDbPath();
     ensureDatabaseInitialized(dbPath);
@@ -146,14 +205,21 @@ export class CloudAccountRepo {
   }
 
   static async migrateToEncrypted(): Promise<void> {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
     try {
-      const rows = db.prepare('SELECT id, token_json, quota_json FROM accounts').all() as any[];
+      const rows = orm
+        .select({
+          id: accounts.id,
+          tokenJson: accounts.tokenJson,
+          quotaJson: accounts.quotaJson,
+        })
+        .from(accounts)
+        .all();
 
       for (const row of rows) {
         let changed = false;
-        let newToken = row.token_json;
-        let newQuota = row.quota_json;
+        let newToken = row.tokenJson;
+        let newQuota = row.quotaJson;
 
         // Check if plain text (starts with {)
         if (newToken && newToken.startsWith('{')) {
@@ -166,105 +232,109 @@ export class CloudAccountRepo {
         }
 
         if (changed) {
-          db.prepare('UPDATE accounts SET token_json = ?, quota_json = ? WHERE id = ?').run(
-            newToken,
-            newQuota,
-            row.id,
-          );
+          orm
+            .update(accounts)
+            .set({ tokenJson: newToken, quotaJson: newQuota })
+            .where(eq(accounts.id, row.id))
+            .run();
           logger.info(`Migrated account ${row.id} to encrypted storage`);
         }
       }
     } catch (e) {
       logger.error('Failed to migrate data', e);
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static async addAccount(account: CloudAccount): Promise<void> {
-    const db = getDb();
-
+    const { raw, orm } = getCloudDb();
     try {
       const tokenEncrypted = await encrypt(JSON.stringify(account.token));
       const quotaEncrypted = account.quota ? await encrypt(JSON.stringify(account.quota)) : null;
+      const values = {
+        id: account.id,
+        provider: account.provider,
+        email: account.email,
+        name: account.name ?? null,
+        avatarUrl: account.avatar_url ?? null,
+        tokenJson: tokenEncrypted,
+        quotaJson: quotaEncrypted,
+        createdAt: account.created_at,
+        lastUsed: account.last_used,
+        status: account.status || 'active',
+        isActive: account.is_active ? 1 : 0,
+      };
 
-      const transaction = db.transaction(() => {
+      orm.transaction((tx) => {
         // If this account is being set to active, deactivate all others first
         if (account.is_active) {
           logger.info(
             `[DEBUG] addAccount: Deactivating all other accounts because ${account.email} is active`,
           );
-          const info = db.prepare('UPDATE accounts SET is_active = 0').run();
+          const info = tx.update(accounts).set({ isActive: 0 }).run();
           logger.info(`[DEBUG] addAccount: Deactivation changed ${info.changes} rows`);
         }
-
-        const stmt = db.prepare(`
-          INSERT OR REPLACE INTO accounts (
-            id, provider, email, name, avatar_url, token_json, quota_json, created_at, last_used, status, is_active
-          ) VALUES (
-            @id, @provider, @email, @name, @avatar_url, @token_json, @quota_json, @created_at, @last_used, @status, @is_active
-          )
-        `);
-
-        stmt.run({
-          id: account.id,
-          provider: account.provider,
-          email: account.email,
-          name: account.name || null,
-          avatar_url: account.avatar_url || null,
-          token_json: tokenEncrypted,
-          quota_json: quotaEncrypted,
-          created_at: account.created_at,
-          last_used: account.last_used,
-          status: account.status || 'active',
-          is_active: account.is_active ? 1 : 0,
-        });
+        tx.insert(accounts)
+          .values(values)
+          .onConflictDoUpdate({
+            target: accounts.id,
+            set: values,
+          })
+          .run();
       });
-
-      transaction();
       logger.info(`Added/Updated cloud account: ${account.email}`);
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static async getAccounts(): Promise<CloudAccount[]> {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
     const migrationStats = createMigrationStats();
 
     try {
-      const stmt = db.prepare('SELECT * FROM accounts ORDER BY last_used DESC');
-      const rows = stmt.all() as any[];
+      const rows = orm.select().from(accounts).orderBy(desc(accounts.lastUsed)).all();
 
       // DEBUG LOGS
-      const activeRows = rows.filter((r) => r.is_active);
+      const activeRows = rows.filter((r) => r.isActive);
       logger.info(
         `[DEBUG] getAccounts: Found ${rows.length} accounts, ${activeRows.length} active.`,
       );
       activeRows.forEach((r) => logger.info(`[DEBUG] Active Account: ${r.email} (${r.id})`));
 
-      const accounts: CloudAccount[] = [];
-      for (const row of rows) {
+      const cloudAccounts: CloudAccount[] = [];
+      for (const normalizedRow of rows) {
         let tokenResult: DecryptFieldResult;
         try {
-          tokenResult = await decryptAndMigrateField(db, row.id, 'token_json', row.token_json);
+          tokenResult = await decryptAndMigrateField(
+            orm,
+            normalizedRow.id,
+            'tokenJson',
+            normalizedRow.tokenJson,
+          );
         } catch (error) {
           migrationStats.failedFields += 1;
-          logger.error(`Failed to decrypt token for account ${row.id}`, error);
+          logger.error(`Failed to decrypt token for account ${normalizedRow.id}`, error);
           throw error;
         }
 
         let quotaResult: DecryptFieldResult;
         try {
-          quotaResult = await decryptAndMigrateField(db, row.id, 'quota_json', row.quota_json);
+          quotaResult = await decryptAndMigrateField(
+            orm,
+            normalizedRow.id,
+            'quotaJson',
+            normalizedRow.quotaJson,
+          );
         } catch (error) {
           migrationStats.failedFields += 1;
-          logger.error(`Failed to decrypt quota for account ${row.id}`, error);
+          logger.error(`Failed to decrypt quota for account ${normalizedRow.id}`, error);
           throw error;
         }
 
         if (!tokenResult.value) {
-          throw new Error(`Missing token data for account ${row.id}`);
+          throw new Error(`Missing token data for account ${normalizedRow.id}`);
         }
 
         if (tokenResult.value) {
@@ -293,22 +363,22 @@ export class CloudAccountRepo {
           }
         }
 
-        accounts.push({
-          id: row.id,
-          provider: row.provider,
-          email: row.email,
-          name: row.name,
-          avatar_url: row.avatar_url,
+        cloudAccounts.push({
+          id: normalizedRow.id,
+          provider: normalizedRow.provider as CloudAccount['provider'],
+          email: normalizedRow.email,
+          name: normalizedRow.name ?? undefined,
+          avatar_url: normalizedRow.avatarUrl ?? undefined,
           token: JSON.parse(tokenResult.value),
           quota: quotaResult.value ? JSON.parse(quotaResult.value) : undefined,
-          created_at: row.created_at,
-          last_used: row.last_used,
-          status: row.status,
-          is_active: Boolean(row.is_active),
+          created_at: normalizedRow.createdAt,
+          last_used: normalizedRow.lastUsed,
+          status: (normalizedRow.status as CloudAccount['status']) ?? undefined,
+          is_active: Boolean(normalizedRow.isActive),
         });
       }
 
-      return accounts;
+      return cloudAccounts;
     } finally {
       if (
         migrationStats.migratedFields > 0 ||
@@ -328,106 +398,321 @@ export class CloudAccountRepo {
           logger.info('CloudAccountRepo migration summary', summary);
         }
       }
-      db.close();
+      raw.close();
     }
   }
 
   static async getAccount(id: string): Promise<CloudAccount | undefined> {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
 
     try {
-      const stmt = db.prepare('SELECT * FROM accounts WHERE id = ?');
-      const row = stmt.get(id) as any;
+      const rows = orm.select().from(accounts).where(eq(accounts.id, id)).all();
+      const normalizedRow = rows[0];
+      if (!normalizedRow) {
+        return undefined;
+      }
 
-      if (!row) return undefined;
-
-      const tokenResult = await decryptAndMigrateField(db, row.id, 'token_json', row.token_json);
-      const quotaResult = await decryptAndMigrateField(db, row.id, 'quota_json', row.quota_json);
+      const tokenResult = await decryptAndMigrateField(
+        orm,
+        normalizedRow.id,
+        'tokenJson',
+        normalizedRow.tokenJson,
+      );
+      const quotaResult = await decryptAndMigrateField(
+        orm,
+        normalizedRow.id,
+        'quotaJson',
+        normalizedRow.quotaJson,
+      );
 
       if (!tokenResult.value) {
-        throw new Error(`Missing token data for account ${row.id}`);
+        throw new Error(`Missing token data for account ${normalizedRow.id}`);
       }
 
       return {
-        id: row.id,
-        provider: row.provider,
-        email: row.email,
-        name: row.name,
-        avatar_url: row.avatar_url,
+        id: normalizedRow.id,
+        provider: normalizedRow.provider as CloudAccount['provider'],
+        email: normalizedRow.email,
+        name: normalizedRow.name ?? undefined,
+        avatar_url: normalizedRow.avatarUrl ?? undefined,
         token: JSON.parse(tokenResult.value),
         quota: quotaResult.value ? JSON.parse(quotaResult.value) : undefined,
-        created_at: row.created_at,
-        last_used: row.last_used,
-        status: row.status,
-        is_active: Boolean(row.is_active),
+        created_at: normalizedRow.createdAt,
+        last_used: normalizedRow.lastUsed,
+        status: (normalizedRow.status as CloudAccount['status']) ?? undefined,
+        is_active: Boolean(normalizedRow.isActive),
       };
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static async removeAccount(id: string): Promise<void> {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
     try {
-      db.prepare('DELETE FROM accounts WHERE id = ?').run(id);
+      orm.delete(accounts).where(eq(accounts.id, id)).run();
       logger.info(`Removed cloud account: ${id}`);
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static async updateToken(id: string, token: any): Promise<void> {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
 
     try {
       const encrypted = await encrypt(JSON.stringify(token));
-      db.prepare('UPDATE accounts SET token_json = ? WHERE id = ?').run(encrypted, id);
+      orm.update(accounts).set({ tokenJson: encrypted }).where(eq(accounts.id, id)).run();
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static async updateQuota(id: string, quota: any): Promise<void> {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
 
     try {
       const encrypted = await encrypt(JSON.stringify(quota));
-      db.prepare('UPDATE accounts SET quota_json = ? WHERE id = ?').run(encrypted, id);
+      orm.update(accounts).set({ quotaJson: encrypted }).where(eq(accounts.id, id)).run();
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static updateLastUsed(id: string): void {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
     try {
-      db.prepare('UPDATE accounts SET last_used = ? WHERE id = ?').run(
-        Math.floor(Date.now() / 1000),
-        id,
-      );
+      orm
+        .update(accounts)
+        .set({ lastUsed: Math.floor(Date.now() / 1000) })
+        .where(eq(accounts.id, id))
+        .run();
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static setActive(id: string): void {
-    const db = getDb();
-    const updateAll = db.prepare('UPDATE accounts SET is_active = 0');
-    const updateOne = db.prepare('UPDATE accounts SET is_active = 1 WHERE id = ?');
-
-    const transaction = db.transaction(() => {
-      updateAll.run();
-      updateOne.run(id);
-    });
+    const { raw, orm } = getCloudDb();
 
     try {
-      transaction();
+      orm.transaction((tx) => {
+        tx.update(accounts).set({ isActive: 0 }).run();
+        tx.update(accounts).set({ isActive: 1 }).where(eq(accounts.id, id)).run();
+      });
       logger.info(`Set account ${id} as active`);
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
+  private static upsertItemValue(db: DrizzleExecutor, key: string, value: string): void {
+    db.insert(itemTable)
+      .values({ key, value })
+      .onConflictDoUpdate({
+        target: itemTable.key,
+        set: { value },
+      })
+      .run();
+  }
+
+  private static writeAuthStatusAndCleanup(db: DrizzleExecutor, account: CloudAccount): void {
+    const authStatus = {
+      name: account.name || account.email,
+      email: account.email,
+      apiKey: account.token.access_token,
+    };
+
+    this.upsertItemValue(db, 'antigravityAuthStatus', JSON.stringify(authStatus));
+    this.upsertItemValue(db, 'antigravityOnboarding', 'true');
+    db.delete(itemTable).where(eq(itemTable.key, 'google.antigravity')).run();
+  }
+
+  private static getItemValue(db: DrizzleExecutor, key: string, context: string): string | null {
+    const rows = db
+      .select({ value: itemTable.value })
+      .from(itemTable)
+      .where(eq(itemTable.key, key))
+      .all();
+    const row = parseRow(ItemTableValueRowSchema, rows[0], context);
+    return row?.value ?? null;
+  }
+
+  private static injectNewFormat(
+    orm: BetterSQLite3Database<typeof drizzleSchema>,
+    account: CloudAccount,
+  ): void {
+    const oauthToken = ProtobufUtils.createUnifiedOAuthToken(
+      account.token.access_token,
+      account.token.refresh_token,
+      account.token.expiry_timestamp,
+    );
+
+    orm.transaction((tx) => {
+      this.upsertItemValue(tx, 'antigravityUnifiedStateSync.oauthToken', oauthToken);
+      this.writeAuthStatusAndCleanup(tx, account);
+    });
+  }
+
+  private static injectOldFormat(
+    orm: BetterSQLite3Database<typeof drizzleSchema>,
+    account: CloudAccount,
+  ): void {
+    const value = this.getItemValue(
+      orm,
+      'jetskiStateSync.agentManagerInitState',
+      'ide.itemTable.jetskiStateSync.agentManagerInitState',
+    );
+
+    orm.transaction((tx) => {
+      if (!value) {
+        logger.warn(
+          'jetskiStateSync.agentManagerInitState not found. ' +
+            'Injecting minimal auth state only. User may need to complete onboarding in the IDE first.',
+        );
+
+        this.writeAuthStatusAndCleanup(tx, account);
+
+        logger.info(
+          `Injected minimal auth state for ${account.email} (no protobuf state available)`,
+        );
+        return;
+      }
+
+      const buffer = Buffer.from(value, 'base64');
+      const data = new Uint8Array(buffer);
+      const cleanData = ProtobufUtils.removeField(data, 6);
+      const newField = ProtobufUtils.createOAuthTokenInfo(
+        account.token.access_token,
+        account.token.refresh_token,
+        account.token.expiry_timestamp,
+      );
+
+      const finalData = new Uint8Array(cleanData.length + newField.length);
+      finalData.set(cleanData, 0);
+      finalData.set(newField, cleanData.length);
+
+      const finalB64 = Buffer.from(finalData).toString('base64');
+
+      tx.update(itemTable)
+        .set({ value: finalB64 })
+        .where(eq(itemTable.key, 'jetskiStateSync.agentManagerInitState'))
+        .run();
+
+      this.writeAuthStatusAndCleanup(tx, account);
+    });
+  }
+
+  private static detectFormatCapability(db: DrizzleExecutor): 'new' | 'old' | null {
+    const unifiedValue = this.getItemValue(
+      db,
+      'antigravityUnifiedStateSync.oauthToken',
+      'ide.itemTable.antigravityUnifiedStateSync.oauthToken',
+    );
+    if (unifiedValue) {
+      return 'new';
+    }
+
+    const oldValue = this.getItemValue(
+      db,
+      'jetskiStateSync.agentManagerInitState',
+      'ide.itemTable.jetskiStateSync.agentManagerInitState',
+    );
+    if (oldValue) {
+      return 'old';
+    }
+
+    return null;
+  }
+
+  private static resolveInjectionStrategy(db: DrizzleExecutor): {
+    name: 'new' | 'old' | 'dual';
+    reason: string;
+  } {
+    try {
+      const version = getAntigravityVersion();
+      return {
+        name: isNewVersion(version) ? 'new' : 'old',
+        reason: `version:${version.shortVersion}`,
+      };
+    } catch (error) {
+      if (!this.versionFailureLogged) {
+        logger.warn('Version detection failed, falling back to capability detection', error);
+        this.versionFailureLogged = true;
+      }
+    }
+
+    const capability = this.detectFormatCapability(db);
+    if (capability) {
+      return { name: capability, reason: 'capability' };
+    }
+
+    return { name: 'dual', reason: 'fallback' };
+  }
+
+  private static getStrategy(name: 'new' | 'old'): {
+    name: 'new' | 'old';
+    inject: (db: BetterSQLite3Database<typeof drizzleSchema>, account: CloudAccount) => void;
+  } {
+    if (name === 'new') {
+      return { name, inject: (db, account) => this.injectNewFormat(db, account) };
+    }
+    return { name, inject: (db, account) => this.injectOldFormat(db, account) };
+  }
+
+  private static injectWithRetry(
+    dbPath: string,
+    account: CloudAccount,
+  ): { strategy: string; attempts: number } {
+    let lastError: unknown;
+    for (let attempt = 1; attempt <= SQLITE_MAX_RETRIES; attempt += 1) {
+      const { raw, orm } = getIdeDb(dbPath, false);
+      try {
+        const { name, reason } = this.resolveInjectionStrategy(orm);
+        if (name === 'dual') {
+          let newInjected = false;
+          let oldInjected = false;
+
+          try {
+            this.injectNewFormat(orm, account);
+            newInjected = true;
+          } catch (newError) {
+            logger.warn('Failed to inject new format', newError);
+          }
+
+          try {
+            this.injectOldFormat(orm, account);
+            oldInjected = true;
+          } catch (oldError) {
+            logger.warn('Failed to inject old format', oldError);
+          }
+
+          if (!newInjected && !oldInjected) {
+            throw new Error('Token injection failed for both formats');
+          }
+
+          return { strategy: `dual:${reason}`, attempts: attempt };
+        }
+
+        const strategy = this.getStrategy(name);
+        strategy.inject(orm, account);
+        return { strategy: `${strategy.name}:${reason}`, attempts: attempt };
+      } catch (error) {
+        lastError = error;
+        if (isSqliteBusyError(error) && attempt < SQLITE_MAX_RETRIES) {
+          logger.warn(`SQLite busy, retrying injection (attempt ${attempt})`, error);
+          sleepSync(SQLITE_RETRY_DELAY_MS);
+          continue;
+        }
+        throw error;
+      } finally {
+        raw.close();
+      }
+    }
+
+    throw lastError;
+  }
+
   static injectCloudToken(account: CloudAccount): void {
     const dbPaths = getAntigravityDbPaths();
     const dbPath = dbPaths.find((p) => fs.existsSync(p)) ?? null;
@@ -436,133 +721,124 @@ export class CloudAccountRepo {
       throw new Error(`Antigravity database not found. Checked paths: ${dbPaths.join(', ')}`);
     }
 
-    const db = new Database(dbPath);
-    db.pragma('journal_mode = WAL');
-    try {
-      const row = db
-        .prepare('SELECT value FROM ItemTable WHERE key = ?')
-        .get('jetskiStateSync.agentManagerInitState') as { value: string } | undefined;
-
-      if (!row || !row.value) {
-        logger.warn(
-          'jetskiStateSync.agentManagerInitState not found. ' +
-            'Injecting minimal auth state only. User may need to complete onboarding in the IDE first.',
-        );
-
-        // Inject minimal state: auth status and onboarding flag only
-        const authStatus = {
-          name: account.name || account.email,
-          email: account.email,
-          apiKey: account.token.access_token,
-        };
-
-        db.prepare('INSERT OR REPLACE INTO ItemTable (key, value) VALUES (?, ?)').run(
-          'antigravityAuthStatus',
-          JSON.stringify(authStatus),
-        );
-
-        db.prepare('INSERT OR REPLACE INTO ItemTable (key, value) VALUES (?, ?)').run(
-          'antigravityOnboarding',
-          'true',
-        );
-
-        db.prepare('DELETE FROM ItemTable WHERE key = ?').run('google.antigravity');
-
-        logger.info(
-          `Injected minimal auth state for ${account.email} (no protobuf state available)`,
-        );
-
-        return; // Early return - skip protobuf manipulation
-      }
-
-      // 1. Decode Base64
-      const buffer = Buffer.from(row.value, 'base64');
-      const data = new Uint8Array(buffer);
-
-      // 2. Remove Field 6
-      const cleanData = ProtobufUtils.removeField(data, 6);
-
-      // 3. Create New Field 6
-      const newField = ProtobufUtils.createOAuthTokenInfo(
-        account.token.access_token,
-        account.token.refresh_token,
-        account.token.expiry_timestamp,
-      );
-
-      // 4. Concatenate
-      const finalData = new Uint8Array(cleanData.length + newField.length);
-      finalData.set(cleanData, 0);
-      finalData.set(newField, cleanData.length);
-
-      // 5. Encode Base64
-      const finalB64 = Buffer.from(finalData).toString('base64');
-
-      // 6. Write back
-      db.prepare('UPDATE ItemTable SET value = ? WHERE key = ?').run(
-        finalB64,
-        'jetskiStateSync.agentManagerInitState',
-      );
-
-      // 7. Inject Onboarding Flag
-      db.prepare('INSERT OR REPLACE INTO ItemTable (key, value) VALUES (?, ?)').run(
-        'antigravityOnboarding',
-        'true',
-      );
-
-      // 8. Update Auth Status (Fix for switching issue)
-      // This ensures the UI reflects the user we just switched to
-      const authStatus = {
-        name: account.name || account.email,
-        email: account.email,
-        apiKey: account.token.access_token, // Critical for session recognition
-        // userStatusProtoBinaryBase64: ... // We cannot generate this easily, hoping IDE fetches it
-      };
-
-      db.prepare('INSERT OR REPLACE INTO ItemTable (key, value) VALUES (?, ?)').run(
-        'antigravityAuthStatus',
-        JSON.stringify(authStatus),
-      );
-
-      // 9. Remove google.antigravity to prevent conflicts
-      // This key often holds old state that might override our injected state
-      db.prepare('DELETE FROM ItemTable WHERE key = ?').run('google.antigravity');
-
-      logger.info(
-        `Successfully injected cloud token and identity for ${account.email} into Antigravity database at ${dbPath}.`,
-      );
-    } finally {
-      db.close();
-    }
+    const result = this.injectWithRetry(dbPath, account);
+    logger.info(
+      `Successfully injected cloud token and identity for ${account.email} into Antigravity database at ${dbPath} (strategy=${result.strategy}, attempts=${result.attempts}).`,
+    );
   }
 
   static getSetting<T>(key: string, defaultValue: T): T {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
     try {
-      const row = db.prepare('SELECT value FROM settings WHERE key = ?').get(key) as
-        | { value: string }
-        | undefined;
-      if (!row) return defaultValue;
+      const rows = orm
+        .select({ value: settings.value })
+        .from(settings)
+        .where(eq(settings.key, key))
+        .all();
+      const row = rows[0];
+      if (!row) {
+        return defaultValue;
+      }
       return JSON.parse(row.value) as T;
     } catch (e) {
       logger.error(`Failed to get setting ${key}`, e);
       return defaultValue;
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
   static setSetting(key: string, value: any): void {
-    const db = getDb();
+    const { raw, orm } = getCloudDb();
     try {
-      db.prepare('INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)').run(
-        key,
-        JSON.stringify(value),
-      );
+      const stringValue = JSON.stringify(value);
+      orm
+        .insert(settings)
+        .values({ key, value: stringValue })
+        .onConflictDoUpdate({
+          target: settings.key,
+          set: { value: stringValue },
+        })
+        .run();
     } finally {
-      db.close();
+      raw.close();
     }
   }
 
+  private static readTokenInfoFromDb(db: DrizzleExecutor): {
+    accessToken: string;
+    refreshToken: string;
+  } {
+    const unifiedValue = this.getItemValue(
+      db,
+      'antigravityUnifiedStateSync.oauthToken',
+      'ide.itemTable.antigravityUnifiedStateSync.oauthToken',
+    );
+
+    let tokenInfo: { accessToken: string; refreshToken: string } | null = null;
+    if (unifiedValue) {
+      try {
+        const unifiedBuffer = Buffer.from(unifiedValue, 'base64');
+        const unifiedData = new Uint8Array(unifiedBuffer);
+        tokenInfo = ProtobufUtils.extractOAuthTokenInfoFromUnifiedState(unifiedData);
+      } catch (error) {
+        logger.warn('SyncLocal: Failed to parse unified OAuth token', error);
+      }
+    }
+
+    if (!tokenInfo) {
+      const value = this.getItemValue(
+        db,
+        'jetskiStateSync.agentManagerInitState',
+        'ide.itemTable.jetskiStateSync.agentManagerInitState',
+      );
+
+      if (!value) {
+        const errorMsg =
+          'No cloud account found in IDE. Please login to a Google account in Antigravity IDE first.';
+        logger.warn(`SyncLocal: ${errorMsg}`);
+        throw new Error(errorMsg);
+      }
+
+      const buffer = Buffer.from(value, 'base64');
+      const data = new Uint8Array(buffer);
+      tokenInfo = ProtobufUtils.extractOAuthTokenInfo(data);
+    }
+
+    if (!tokenInfo) {
+      const errorMsg =
+        'No OAuth token found in IDE state. Please login to a Google account in Antigravity IDE first.';
+      logger.warn(`SyncLocal: ${errorMsg}`);
+      throw new Error(errorMsg);
+    }
+
+    return tokenInfo;
+  }
+
+  private static readTokenInfoWithRetry(dbPath: string): {
+    accessToken: string;
+    refreshToken: string;
+  } {
+    let lastError: unknown;
+    for (let attempt = 1; attempt <= SQLITE_MAX_RETRIES; attempt += 1) {
+      const { raw, orm } = getIdeDb(dbPath, true);
+      try {
+        return this.readTokenInfoFromDb(orm);
+      } catch (error) {
+        lastError = error;
+        if (isSqliteBusyError(error) && attempt < SQLITE_MAX_RETRIES) {
+          logger.warn(`SQLite busy, retrying IDE read (attempt ${attempt})`, error);
+          sleepSync(SQLITE_RETRY_DELAY_MS);
+          continue;
+        }
+        throw error;
+      } finally {
+        raw.close();
+      }
+    }
+    throw lastError;
+  }
+
   static async syncFromIDE(): Promise<CloudAccount | null> {
     // Try all possible database paths
     const dbPaths = getAntigravityDbPaths();
@@ -581,32 +857,8 @@ export class CloudAccountRepo {
     }
 
     logger.info(`SyncLocal: Using Antigravity database at: ${dbPath}`);
-    const ideDb = new Database(dbPath, { readonly: true });
-    ideDb.pragma('journal_mode = WAL');
     try {
-      // 1. Read Raw Token Data
-      const row = ideDb
-        .prepare('SELECT value FROM ItemTable WHERE key = ?')
-        .get('jetskiStateSync.agentManagerInitState') as { value: string } | undefined;
-
-      if (!row || !row.value) {
-        const errorMsg =
-          'No cloud account found in IDE. Please login to a Google account in Antigravity IDE first.';
-        logger.warn(`SyncLocal: ${errorMsg}`);
-        throw new Error(errorMsg);
-      }
-
-      // 2. Decode Protobuf
-      const buffer = Buffer.from(row.value, 'base64');
-      const data = new Uint8Array(buffer);
-      const tokenInfo = ProtobufUtils.extractOAuthTokenInfo(data);
-
-      if (!tokenInfo) {
-        const errorMsg =
-          'No OAuth token found in IDE state. Please login to a Google account in Antigravity IDE first.';
-        logger.warn(`SyncLocal: ${errorMsg}`);
-        throw new Error(errorMsg);
-      }
+      const tokenInfo = this.readTokenInfoWithRetry(dbPath);
 
       // 3. Fetch User Info
       // We need to fetch user info to know who this token belongs to
@@ -656,8 +908,6 @@ export class CloudAccountRepo {
     } catch (error) {
       logger.error('SyncLocal: Failed to sync account from IDE', error);
       throw error;
-    } finally {
-      ideDb.close();
     }
   }
 }