evandro/gemini-nexus
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
vin authored and GitHub committed 2026-02-06 20:16:02 +08:00
1 parent f572ae4652
commit ed94abf0d6
16 files changed
+1658 -328
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/ipc/database/cloudHandler.ts
+508 -258
View File
@@ -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();
}
}
}
src/ipc/database/dbConnection.ts
+34
View File
@@ -0,0 +1,34 @@
import Database from 'better-sqlite3';
import { drizzle, type BetterSQLite3Database } from 'drizzle-orm/better-sqlite3';
import * as schema from './schema';
export interface DbConfigOptions {
readOnly?: boolean;
busyTimeoutMs?: number;
enableWal?: boolean;
}
export interface DrizzleConnection {
raw: Database.Database;
orm: BetterSQLite3Database<typeof schema>;
}
export function configureDatabase(db: Database.Database, options: DbConfigOptions = {}): void {
const busyTimeoutMs = options.busyTimeoutMs ?? 3000;
const enableWal = options.enableWal ?? !options.readOnly;
if (enableWal) {
db.pragma('journal_mode = WAL');
}
db.pragma(`busy_timeout = ${busyTimeoutMs}`);
}
export function openDrizzleConnection(
dbPath: string,
options: Database.Options = {},
config: DbConfigOptions = {},
): DrizzleConnection {
const raw = new Database(dbPath, options);
configureDatabase(raw, config);
const orm = drizzle(raw, { schema });
return { raw, orm };
}
src/ipc/database/handler.ts
+95 -66
View File
@@ -1,17 +1,27 @@
import Database from 'better-sqlite3';
import fs from 'fs';
import { getAntigravityDbPaths } from '../../utils/paths';
import { logger } from '../../utils/logger';
import { AccountBackupData, AccountInfo } from '../../types/account';
import path from 'path';
import { eq } from 'drizzle-orm';
import { AccountBackupData, AccountInfo } from '../../types/account';
import { ItemTableValueRowSchema, type ItemTableKey } from '../../types/db';
import { logger } from '../../utils/logger';
import { getAntigravityDbPaths } from '../../utils/paths';
import { parseRow } from '../../utils/sqlite';
import { openDrizzleConnection } from './dbConnection';
import { itemTable } from './schema';
const KEYS_TO_BACKUP = ['antigravityAuthStatus', 'jetskiStateSync.agentManagerInitState'];
const KEYS_TO_BACKUP: ItemTableKey[] = [
'antigravityAuthStatus',
'jetskiStateSync.agentManagerInitState',
'antigravityUnifiedStateSync.oauthToken',
];
function getDbWithWal(dbPath: string, options: Database.Options = {}): Database.Database {
const db = new Database(dbPath, options);
db.pragma('journal_mode = WAL');
return db;
function openIdeDb(dbPath: string, readOnly = false): ReturnType<typeof openDrizzleConnection> {
return openDrizzleConnection(
dbPath,
{ readonly: readOnly, fileMustExist: false },
{ readOnly, busyTimeoutMs: 3000 },
);
}
/**
@@ -25,8 +35,8 @@ export function initDatabase(): void {
return;
}
const db = getDatabaseConnection();
db.close();
const { raw } = getDatabaseConnection();
raw.close();
logger.info('Database initialized and verified (WAL mode)');
} catch (error) {
logger.error('Failed to initialize database on startup', error);
@@ -72,9 +82,9 @@ function ensureDatabaseExists(dbPath: string): void {
/**
* Gets a database connection.
* @param dbPath {string} The path to the database file.
* @returns {Database.Database} The database connection.
* @returns {ReturnType<typeof openDrizzleConnection>} The database connection.
*/
export function getDatabaseConnection(dbPath?: string): Database.Database {
export function getDatabaseConnection(dbPath?: string): ReturnType<typeof openDrizzleConnection> {
const targetPath = dbPath || getAntigravityDbPaths()[0];
if (!targetPath) {
@@ -84,7 +94,7 @@ export function getDatabaseConnection(dbPath?: string): Database.Database {
ensureDatabaseExists(targetPath);
try {
return getDbWithWal(targetPath, { readonly: false, fileMustExist: false });
return openIdeDb(targetPath);
} catch (error: unknown) {
const err = error as { code?: string };
if (err.code === 'SQLITE_BUSY' || err.code === 'SQLITE_LOCKED') {
@@ -94,11 +104,19 @@ export function getDatabaseConnection(dbPath?: string): Database.Database {
}
}
/**
* Gets a database connection and enables WAL mode.
* @param dbPath {string} The path to the database file.
* @returns {Database.Database} The database connection.
*/
function readItemValue(
orm: ReturnType<typeof openDrizzleConnection>['orm'],
key: string,
context: string,
): string | null {
const rows = orm
.select({ value: itemTable.value })
.from(itemTable)
.where(eq(itemTable.key, key))
.all();
const row = parseRow(ItemTableValueRowSchema, rows[0], context);
return row?.value ?? null;
}
/**
* Gets the current account info.
@@ -106,61 +124,66 @@ export function getDatabaseConnection(dbPath?: string): Database.Database {
*/
export function getCurrentAccountInfo(): AccountInfo {
// NOTE Database existence is now handled by getDatabaseConnection
let db: Database.Database | null = null;
let connection: ReturnType<typeof openDrizzleConnection> | null = null;
try {
db = getDatabaseConnection();
connection = getDatabaseConnection();
const { orm } = connection;
// Query for auth status
const authRow = db
.prepare("SELECT value FROM ItemTable WHERE key = 'antigravityAuthStatus'")
.get() as { value: string } | undefined;
const authValue = readItemValue(
orm,
'antigravityAuthStatus',
'ide.itemTable.antigravityAuthStatus',
);
let authStatus = null;
if (authRow) {
if (authValue) {
try {
authStatus = JSON.parse(authRow.value);
authStatus = JSON.parse(authValue);
} catch {
// NOTE Ignore JSON parse errors
}
}
// NOTE Query for user info (usually in jetskiStateSync.agentManagerInitState or similar)
const initRow = db
.prepare("SELECT value FROM ItemTable WHERE key = 'jetskiStateSync.agentManagerInitState'")
.get() as { value: string } | undefined;
const initValue = readItemValue(
orm,
'jetskiStateSync.agentManagerInitState',
'ide.itemTable.jetskiStateSync.agentManagerInitState',
);
let initState = null;
if (initRow) {
if (initValue) {
try {
initState = JSON.parse(initRow.value);
initState = JSON.parse(initValue);
} catch {
// Ignore JSON parse errors (this key often contains non-JSON data)
}
}
// Query for google.antigravity
const googleRow = db
.prepare("SELECT value FROM ItemTable WHERE key = 'google.antigravity'")
.get() as { value: string } | undefined;
const googleValue = readItemValue(
orm,
'google.antigravity',
'ide.itemTable.google.antigravity',
);
let googleState = null;
if (googleRow) {
if (googleValue) {
try {
googleState = JSON.parse(googleRow.value);
googleState = JSON.parse(googleValue);
} catch {
// Ignore JSON parse errors
}
}
// Query for antigravityUserSettings.allUserSettings
const settingsRow = db
.prepare("SELECT value FROM ItemTable WHERE key = 'antigravityUserSettings.allUserSettings'")
.get() as { value: string } | undefined;
const settingsValue = readItemValue(
orm,
'antigravityUserSettings.allUserSettings',
'ide.itemTable.antigravityUserSettings.allUserSettings',
);
let settingsState = null;
if (settingsRow) {
if (settingsValue) {
try {
settingsState = JSON.parse(settingsRow.value);
settingsState = JSON.parse(settingsValue);
} catch {
// Ignore JSON parse errors
}
@@ -195,28 +218,28 @@ export function getCurrentAccountInfo(): AccountInfo {
logger.error('Failed to get current account info', error);
throw error;
} finally {
if (db) db.close();
if (connection) {
connection.raw.close();
}
}
}
export function backupAccount(account: AccountBackupData['account']): AccountBackupData {
let db: Database.Database | null = null;
let connection: ReturnType<typeof openDrizzleConnection> | null = null;
try {
db = getDatabaseConnection();
connection = getDatabaseConnection();
const { orm } = connection;
// NOTE Backup only specific keys
const data: Record<string, unknown> = {};
for (const key of KEYS_TO_BACKUP) {
const row = db.prepare('SELECT value FROM ItemTable WHERE key = ?').get(key) as
| { value: string }
| undefined;
if (row) {
const value = readItemValue(orm, key, `ide.itemTable.backup.${key}`);
if (value) {
try {
data[key] = JSON.parse(row.value);
data[key] = JSON.parse(value);
} catch {
data[key] = row.value;
data[key] = value;
}
logger.debug(`Backed up key: ${key}`);
} else {
@@ -237,7 +260,9 @@ export function backupAccount(account: AccountBackupData['account']): AccountBac
logger.error('Failed to backup account', error);
throw error;
} finally {
if (db) db.close();
if (connection) {
connection.raw.close();
}
}
}
@@ -288,32 +313,36 @@ function _restoreSingleDb(dbPath: string, backup: AccountBackupData): boolean {
}
logger.info(`Restoring database: ${dbPath}`);
let db: Database.Database | null = null;
let connection: ReturnType<typeof openDrizzleConnection> | null = null;
try {
db = getDatabaseConnection(dbPath);
const insert = db.prepare('INSERT OR REPLACE INTO ItemTable (key, value) VALUES (?, ?)');
const transaction = db.transaction(() => {
connection = getDatabaseConnection(dbPath);
const { orm } = connection;
orm.transaction((tx) => {
// NOTE Only restore the keys that were backed up
for (const key of KEYS_TO_BACKUP) {
if (key in backup.data) {
const value = backup.data[key];
const stringValue = typeof value === 'string' ? value : JSON.stringify(value);
insert.run(key, stringValue);
tx.insert(itemTable)
.values({ key, value: stringValue })
.onConflictDoUpdate({
target: itemTable.key,
set: { value: stringValue },
})
.run();
logger.debug(`Restored key: ${key}`);
}
}
});
transaction();
logger.info(`Database restoration complete: ${dbPath}`);
return true;
} catch (error) {
logger.error(`Failed to restore database: ${dbPath}`, error);
return false;
} finally {
if (db) db.close();
if (connection) {
connection.raw.close();
}
}
}
src/ipc/database/schema.ts
+25
View File
@@ -0,0 +1,25 @@
import { integer, sqliteTable, text } from 'drizzle-orm/sqlite-core';
export const accounts = sqliteTable('accounts', {
id: text('id').primaryKey(),
provider: text('provider').notNull(),
email: text('email').notNull(),
name: text('name'),
avatarUrl: text('avatar_url'),
tokenJson: text('token_json').notNull(),
quotaJson: text('quota_json'),
createdAt: integer('created_at').notNull(),
lastUsed: integer('last_used').notNull(),
status: text('status'),
isActive: integer('is_active').notNull().default(0),
});
export const settings = sqliteTable('settings', {
key: text('key').primaryKey(),
value: text('value').notNull(),
});
export const itemTable = sqliteTable('ItemTable', {
key: text('key').primaryKey(),
value: text('value'),
});