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

src/tests/unit/antigravityVersion.test.ts

@@ -0,0 +1,57 @@
+import { describe, it, expect, vi } from 'vitest';
+import { compareVersion, isNewVersion } from '../../utils/antigravityVersion';
+
+describe('antigravityVersion', () => {
+  it('should compare versions correctly', () => {
+    expect(compareVersion('1.16.5', '1.16.4')).toBe(1);
+    expect(compareVersion('1.16.5', '1.16.5')).toBe(0);
+    expect(compareVersion('1.16.4', '1.16.5')).toBe(-1);
+    expect(compareVersion('1.17.0', '1.16.5')).toBe(1);
+    expect(compareVersion('2.0.0', '1.16.5')).toBe(1);
+  });
+
+  it('should detect new version >= 1.16.5', () => {
+    expect(isNewVersion({ shortVersion: '1.16.4', bundleVersion: '1.16.4' })).toBe(false);
+    expect(isNewVersion({ shortVersion: '1.16.5', bundleVersion: '1.16.5' })).toBe(true);
+    expect(isNewVersion({ shortVersion: '1.17.0', bundleVersion: '1.17.0' })).toBe(true);
+  });
+
+  it('should fallback to package.json on windows when version detection fails', async () => {
+    if (process.platform !== 'win32') {
+      return;
+    }
+
+    vi.resetModules();
+    const execSync = vi.fn(() => {
+      throw new Error('ps fail');
+    });
+    const existsSync = vi.fn(() => true);
+    const readFileSync = vi.fn(() => JSON.stringify({ version: '1.16.6' }));
+
+    vi.doMock('child_process', () => ({
+      execSync,
+      default: {
+        execSync,
+      },
+    }));
+    vi.doMock('fs', () => ({
+      existsSync,
+      readFileSync,
+      default: {
+        existsSync,
+        readFileSync,
+      },
+    }));
+    vi.doMock('../../utils/paths', () => ({
+      getAntigravityExecutablePath: () => 'C:\\Program Files\\Antigravity\\Antigravity.exe',
+    }));
+
+    const { getAntigravityVersion } = await import('../../utils/antigravityVersion');
+    const version = getAntigravityVersion();
+    const cached = getAntigravityVersion();
+
+    expect(version.shortVersion).toBe('1.16.6');
+    expect(cached.shortVersion).toBe('1.16.6');
+    expect(execSync).toHaveBeenCalledTimes(1);
+  });
+});

src/tests/unit/cloudHandler-sync.test.ts

@@ -0,0 +1,269 @@
+import fs from 'fs';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { ProtobufUtils } from '../../utils/protobuf';
+import { CloudAccountRepo } from '../../ipc/database/cloudHandler';
+import type { UserInfo } from '../../services/GoogleAPIService';
+
+let mockData: Record<string, string>;
+let busyOnFirstGet = false;
+let getCallCount = 0;
+let runCalls: Array<{ sql: string; args: unknown[] }>;
+interface MockOrm {
+  select: () => {
+    from: () => {
+      where: (condition: { __key?: string }) => { all: () => Array<{ value: string }> };
+    };
+  };
+  insert: () => {
+    values: (values: { key: string; value: string }) => {
+      onConflictDoUpdate: () => { run: () => { changes: number } };
+    };
+  };
+  update: () => {
+    set: (values: { value?: string }) => {
+      where: (condition: { __key?: string }) => { run: () => { changes: number } };
+    };
+  };
+  delete: () => {
+    where: (condition: { __key?: string }) => { run: () => { changes: number } };
+  };
+  transaction: (fn: (tx: MockOrm) => void) => void;
+}
+
+let mockOrm: MockOrm;
+
+function createMockUserInfo(email: string, name: string): UserInfo {
+  return {
+    id: `id-${email}`,
+    email,
+    verified_email: true,
+    name,
+    given_name: name,
+    family_name: 'User',
+    picture: '',
+  };
+}
+
+vi.mock('drizzle-orm', () => ({
+  eq: (_column: unknown, value: string) => ({ __key: value }),
+  desc: (value: unknown) => value,
+}));
+
+vi.mock('../../ipc/database/dbConnection', () => ({
+  openDrizzleConnection: () => ({
+    raw: { close: vi.fn() },
+    orm: mockOrm,
+  }),
+}));
+
+vi.mock('../../utils/paths', () => ({
+  getAntigravityDbPaths: () => ['mock-db'],
+  getCloudAccountsDbPath: () => 'mock-cloud-db',
+}));
+
+vi.mock('../../utils/logger', () => ({
+  logger: {
+    info: vi.fn(),
+    warn: vi.fn(),
+    error: vi.fn(),
+    debug: vi.fn(),
+  },
+}));
+
+vi.mock('../../services/GoogleAPIService', () => ({
+  GoogleAPIService: {
+    getUserInfo: vi.fn(),
+  },
+}));
+
+describe('CloudAccountRepo.syncFromIDE', () => {
+  beforeEach(() => {
+    mockData = {};
+    busyOnFirstGet = false;
+    getCallCount = 0;
+    runCalls = [];
+    vi.spyOn(fs, 'existsSync').mockReturnValue(true);
+    mockOrm = {
+      select: () => ({
+        from: () => ({
+          where: (condition: { __key?: string }) => ({
+            all: () => {
+              getCallCount += 1;
+              if (busyOnFirstGet && getCallCount === 1) {
+                const error = new Error('SQLITE_BUSY');
+                (error as { code?: string }).code = 'SQLITE_BUSY';
+                throw error;
+              }
+              const key = condition?.__key ?? '';
+              const value = mockData[key];
+              if (value === undefined) {
+                return [];
+              }
+              return [{ value }];
+            },
+          }),
+        }),
+      }),
+      insert: () => ({
+        values: (values: { key: string; value: string }) => ({
+          onConflictDoUpdate: () => ({
+            run: () => {
+              runCalls.push({ sql: 'insert', args: [values] });
+              return { changes: 1 };
+            },
+          }),
+        }),
+      }),
+      update: () => ({
+        set: (values: { value?: string }) => ({
+          where: (condition: { __key?: string }) => ({
+            run: () => {
+              runCalls.push({ sql: 'update', args: [values, condition] });
+              return { changes: 1 };
+            },
+          }),
+        }),
+      }),
+      delete: () => ({
+        where: (condition: { __key?: string }) => ({
+          run: () => {
+            runCalls.push({ sql: 'delete', args: [condition] });
+            return { changes: 1 };
+          },
+        }),
+      }),
+      transaction: (fn: (tx: typeof mockOrm) => void) => {
+        fn(mockOrm);
+      },
+    };
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('should prefer unified oauth token when present', async () => {
+    const accessToken = 'access-new';
+    const refreshToken = 'refresh-new';
+    const unifiedB64 = ProtobufUtils.createUnifiedOAuthToken(accessToken, refreshToken, 1700000000);
+
+    const oldB64 = Buffer.from(
+      ProtobufUtils.createOAuthTokenInfo('access-old', 'refresh-old', 1700000000),
+    ).toString('base64');
+
+    mockData['antigravityUnifiedStateSync.oauthToken'] = unifiedB64;
+    mockData['jetskiStateSync.agentManagerInitState'] = oldB64;
+
+    const { GoogleAPIService } = await import('../../services/GoogleAPIService');
+    vi.mocked(GoogleAPIService.getUserInfo).mockResolvedValue(
+      createMockUserInfo('[email protected]', 'New User'),
+    );
+
+    vi.spyOn(CloudAccountRepo, 'getAccounts').mockResolvedValue([]);
+    vi.spyOn(CloudAccountRepo, 'addAccount').mockResolvedValue();
+
+    const account = await CloudAccountRepo.syncFromIDE();
+
+    expect(GoogleAPIService.getUserInfo).toHaveBeenCalledWith(accessToken);
+    expect(account?.email).toBe('[email protected]');
+  });
+
+  it('should fall back to old oauth token when unified is missing', async () => {
+    const accessToken = 'access-old';
+    const refreshToken = 'refresh-old';
+    const oldB64 = Buffer.from(
+      ProtobufUtils.createOAuthTokenInfo(accessToken, refreshToken, 1700000000),
+    ).toString('base64');
+
+    mockData['jetskiStateSync.agentManagerInitState'] = oldB64;
+
+    const { GoogleAPIService } = await import('../../services/GoogleAPIService');
+    vi.mocked(GoogleAPIService.getUserInfo).mockResolvedValue(
+      createMockUserInfo('[email protected]', 'Old User'),
+    );
+
+    vi.spyOn(CloudAccountRepo, 'getAccounts').mockResolvedValue([]);
+    vi.spyOn(CloudAccountRepo, 'addAccount').mockResolvedValue();
+
+    const account = await CloudAccountRepo.syncFromIDE();
+
+    expect(GoogleAPIService.getUserInfo).toHaveBeenCalledWith(accessToken);
+    expect(account?.email).toBe('[email protected]');
+  });
+
+  it('should retry when sqlite is busy', async () => {
+    busyOnFirstGet = true;
+    const accessToken = 'access-retry';
+    const refreshToken = 'refresh-retry';
+    const oldB64 = Buffer.from(
+      ProtobufUtils.createOAuthTokenInfo(accessToken, refreshToken, 1700000000),
+    ).toString('base64');
+
+    mockData['jetskiStateSync.agentManagerInitState'] = oldB64;
+
+    const { GoogleAPIService } = await import('../../services/GoogleAPIService');
+    vi.mocked(GoogleAPIService.getUserInfo).mockResolvedValue(
+      createMockUserInfo('[email protected]', 'Retry User'),
+    );
+
+    vi.spyOn(CloudAccountRepo, 'getAccounts').mockResolvedValue([]);
+    vi.spyOn(CloudAccountRepo, 'addAccount').mockResolvedValue();
+
+    const account = await CloudAccountRepo.syncFromIDE();
+
+    expect(GoogleAPIService.getUserInfo).toHaveBeenCalledWith(accessToken);
+    expect(account?.email).toBe('[email protected]');
+  });
+
+  it('should prefer new format when capability detection finds unified key', async () => {
+    vi.resetModules();
+    vi.doMock('../../utils/antigravityVersion', () => ({
+      getAntigravityVersion: () => {
+        throw new Error('version detection failed');
+      },
+      isNewVersion: () => false,
+    }));
+
+    const { CloudAccountRepo: RepoWithMock } = await import('../../ipc/database/cloudHandler');
+    const accessToken = 'access-new';
+    const refreshToken = 'refresh-new';
+
+    mockData['antigravityUnifiedStateSync.oauthToken'] = 'exists';
+    mockData['jetskiStateSync.agentManagerInitState'] = 'exists-old';
+
+    RepoWithMock.injectCloudToken({
+      id: 'id',
+      provider: 'google',
+      email: '[email protected]',
+      name: 'Test',
+      avatar_url: '',
+      token: {
+        access_token: accessToken,
+        refresh_token: refreshToken,
+        expires_in: 3600,
+        expiry_timestamp: 1700000000,
+        token_type: 'Bearer',
+        email: '[email protected]',
+      },
+      created_at: 1700000000,
+      last_used: 1700000000,
+      status: 'active',
+      is_active: true,
+    });
+
+    const updatedOldKey = runCalls.some(
+      (call) =>
+        call.sql === 'update' &&
+        (call.args[1] as { __key?: string } | undefined)?.__key ===
+          'jetskiStateSync.agentManagerInitState',
+    );
+    const wroteUnifiedKey = runCalls.some(
+      (call) =>
+        call.sql === 'insert' &&
+        (call.args[0] as { key?: string })?.key === 'antigravityUnifiedStateSync.oauthToken',
+    );
+
+    expect(wroteUnifiedKey).toBe(true);
+    expect(updatedOldKey).toBe(false);
+  });
+});

src/tests/unit/database-backup-keys.test.ts

@@ -0,0 +1,84 @@
+import fs from 'fs';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import * as handler from '../../ipc/database/handler';
+
+let writes: Record<string, string> = {};
+interface MockOrm {
+  insert: () => {
+    values: (values: { key: string; value: string }) => {
+      onConflictDoUpdate: () => { run: () => { changes: number } };
+    };
+  };
+  transaction: (fn: (tx: MockOrm) => void) => void;
+}
+
+let mockOrm: MockOrm;
+
+vi.mock('../../ipc/database/dbConnection', () => ({
+  openDrizzleConnection: () => ({
+    raw: { close: vi.fn() },
+    orm: mockOrm,
+  }),
+}));
+
+vi.mock('../../utils/paths', () => ({
+  getAntigravityDbPaths: () => ['mock-db'],
+}));
+
+vi.mock('../../utils/logger', () => ({
+  logger: {
+    info: vi.fn(),
+    warn: vi.fn(),
+    error: vi.fn(),
+    debug: vi.fn(),
+  },
+}));
+
+describe('database backup keys', () => {
+  beforeEach(() => {
+    writes = {};
+    mockOrm = {
+      insert: () => ({
+        values: (values: { key: string; value: string }) => ({
+          onConflictDoUpdate: () => ({
+            run: () => {
+              writes[values.key] = values.value;
+              return { changes: 1 };
+            },
+          }),
+        }),
+      }),
+      transaction: (fn: (tx: typeof mockOrm) => void) => {
+        fn(mockOrm);
+      },
+    };
+    vi.spyOn(fs, 'existsSync').mockImplementation((value) => {
+      const pathValue = String(value);
+      return !pathValue.endsWith('.backup');
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('should restore unified oauth token key', () => {
+    handler.restoreAccount({
+      version: '1.0',
+      account: {
+        id: '1',
+        name: 'Test',
+        email: '[email protected]',
+        created_at: new Date().toISOString(),
+        last_used: new Date().toISOString(),
+      },
+      data: {
+        antigravityAuthStatus: '{"email":"[email protected]"}',
+        'jetskiStateSync.agentManagerInitState': 'old',
+        'antigravityUnifiedStateSync.oauthToken': 'unified',
+      },
+    });
+
+    expect(writes['antigravityUnifiedStateSync.oauthToken']).toBe('unified');
+  });
+});

src/tests/unit/database.test.ts

@@ -54,9 +54,9 @@ describe.skip('Database Handler', () => {
   });
 
   it('should connect to database', () => {
-    const db = getDatabaseConnection();
-    expect(db).toBeDefined();
-    db.close();
+    const { raw } = getDatabaseConnection();
+    expect(raw).toBeDefined();
+    raw.close();
   });
 
   it('should get current account info', () => {

src/tests/unit/protobuf-unified.test.ts

@@ -0,0 +1,37 @@
+import { describe, it, expect } from 'vitest';
+import { ProtobufUtils } from '../../utils/protobuf';
+
+describe('ProtobufUtils Unified OAuth', () => {
+  it('should round-trip OAuthInfo payload', () => {
+    const accessToken = 'access-token-123';
+    const refreshToken = 'refresh-token-456';
+    const expiry = 1700000000;
+
+    const oauthInfo = ProtobufUtils.createOAuthInfo(accessToken, refreshToken, expiry);
+    const parsed = ProtobufUtils.extractOAuthTokenInfoFromOAuthInfo(oauthInfo);
+
+    expect(parsed).toEqual({
+      accessToken,
+      refreshToken,
+    });
+  });
+
+  it('should round-trip unified oauth token', () => {
+    const accessToken = 'access-token-abc';
+    const refreshToken = 'refresh-token-def';
+    const expiry = 1700001234;
+
+    const unifiedB64 = ProtobufUtils.createUnifiedOAuthToken(
+      accessToken,
+      refreshToken,
+      expiry,
+    );
+    const unifiedBytes = new Uint8Array(Buffer.from(unifiedB64, 'base64'));
+    const parsed = ProtobufUtils.extractOAuthTokenInfoFromUnifiedState(unifiedBytes);
+
+    expect(parsed).toEqual({
+      accessToken,
+      refreshToken,
+    });
+  });
+});