Initial Commit
This commit is contained in:
Draculabo committed
commit
d0417339c1
158 files changed
+31368
No files matched your search
@@ -0,0 +1,108 @@
|
||||
import { CloudAccountRepo } from '../ipc/database/cloudHandler';
|
||||
import { CloudAccount } from '../types/cloudAccount';
|
||||
import { switchCloudAccount } from '../ipc/cloud/handler';
|
||||
import { logger } from '../utils/logger';
|
||||
|
||||
export class AutoSwitchService {
|
||||
/**
|
||||
* Finds the best cloud account to switch to.
|
||||
* Criteria:
|
||||
* 1. Not the current account (unless it's the only one).
|
||||
* 2. Status is 'active'.
|
||||
* 3. Has quota > 5% for all models (or at least gemini-pro).
|
||||
* 4. Sorted by highest quota then last_used (least recently used preferred for rotation? or most? Let's say highest quota first).
|
||||
*/
|
||||
static async findBestAccount(currentAccountId: string): Promise<CloudAccount | null> {
|
||||
const accounts = await CloudAccountRepo.getAccounts();
|
||||
|
||||
// Filter potential candidates
|
||||
const candidates = accounts.filter((acc) => {
|
||||
if (acc.id === currentAccountId) return false;
|
||||
if (acc.status !== 'active') return false; // Rate limited or expired accounts are skipped
|
||||
|
||||
// Check quota
|
||||
// We assume simple check: if any model has < 5%, we skip it.
|
||||
// Or better: check average? NO, check critical models.
|
||||
// For now, let's just check if quota object exists.
|
||||
if (!acc.quota) return false; // No quota data means risky
|
||||
|
||||
const models = Object.values(acc.quota.models);
|
||||
// If any model is depleted (< 5%), skip.
|
||||
const isDepleted = models.some((m) => m.percentage < 5);
|
||||
return !isDepleted;
|
||||
});
|
||||
|
||||
if (candidates.length === 0) return null;
|
||||
|
||||
// Sort by "Best"
|
||||
// Heuristic: Highest average quota availability
|
||||
candidates.sort((a, b) => {
|
||||
const avgA = this.calculateAverageQuota(a);
|
||||
const avgB = this.calculateAverageQuota(b);
|
||||
return avgB - avgA; // Descending
|
||||
});
|
||||
|
||||
return candidates[0];
|
||||
}
|
||||
|
||||
private static calculateAverageQuota(account: CloudAccount): number {
|
||||
if (!account.quota) return 0;
|
||||
const values = Object.values(account.quota.models).map((m) => m.percentage);
|
||||
if (values.length === 0) return 0;
|
||||
const sum = values.reduce((a, b) => a + b, 0);
|
||||
return sum / values.length;
|
||||
}
|
||||
|
||||
/**
|
||||
* Triggered by Monitor Service or UI.
|
||||
* Checks if we need to switch from the current account.
|
||||
*/
|
||||
static async checkAndSwitchIfNeeded(): Promise<boolean> {
|
||||
const enabled = CloudAccountRepo.getSetting<boolean>('auto_switch_enabled', false);
|
||||
if (!enabled) return false;
|
||||
|
||||
// Get current active account
|
||||
const accounts = await CloudAccountRepo.getAccounts();
|
||||
const currentAccount = accounts.find((a) => a.is_active);
|
||||
|
||||
// If no active account, maybe we should pick one?
|
||||
// For now, assume user manually picked first one.
|
||||
if (!currentAccount) return false;
|
||||
|
||||
// Check if current is depleted
|
||||
const isDepleted = this.isAccountDepleted(currentAccount);
|
||||
|
||||
if (isDepleted || currentAccount.status === 'rate_limited') {
|
||||
logger.info(
|
||||
`AutoSwitch: Current account ${currentAccount.email} is depleted or rate limited.`,
|
||||
);
|
||||
|
||||
const nextAccount = await this.findBestAccount(currentAccount.id);
|
||||
if (nextAccount) {
|
||||
logger.info(`AutoSwitch: Switching to ${nextAccount.email}...`);
|
||||
|
||||
// Notify user (via toast? we are in main process... IPC event?)
|
||||
// Ideally we send an IPC event to renderer.
|
||||
// For now, logic first.
|
||||
|
||||
await switchCloudAccount(nextAccount.id);
|
||||
|
||||
// We might want to send a notification to user desktop?
|
||||
// require('electron').Notification ...
|
||||
|
||||
return true;
|
||||
} else {
|
||||
logger.warn('AutoSwitch: No healthy accounts available to switch to.');
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
static isAccountDepleted(account: CloudAccount): boolean {
|
||||
if (!account.quota) return false; // Unknown, assume fine or let fetchQuota find out
|
||||
// Threshold = 5%
|
||||
const THRESHOLD = 5;
|
||||
return Object.values(account.quota.models).some((m) => m.percentage < THRESHOLD);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
import { CloudAccountRepo } from '../ipc/database/cloudHandler';
|
||||
import { GoogleAPIService } from './GoogleAPIService';
|
||||
import { AutoSwitchService } from './AutoSwitchService';
|
||||
import { logger } from '../utils/logger';
|
||||
|
||||
export class CloudMonitorService {
|
||||
private static intervalId: NodeJS.Timeout | null = null;
|
||||
private static POLL_INTERVAL = 1000 * 60 * 5; // 5 minutes
|
||||
|
||||
static start() {
|
||||
if (this.intervalId) return;
|
||||
logger.info('Starting CloudMonitorService...');
|
||||
|
||||
// Initial Poll
|
||||
this.poll().catch((e) => logger.error('Initial poll failed', e));
|
||||
|
||||
this.intervalId = setInterval(() => {
|
||||
this.poll().catch((e) => logger.error('Scheduled poll failed', e));
|
||||
}, this.POLL_INTERVAL);
|
||||
}
|
||||
|
||||
static stop() {
|
||||
if (this.intervalId) {
|
||||
clearInterval(this.intervalId);
|
||||
this.intervalId = null;
|
||||
logger.info('Stopped CloudMonitorService');
|
||||
}
|
||||
}
|
||||
|
||||
static async poll() {
|
||||
logger.info('CloudMonitor: Polling quotas...');
|
||||
const accounts = await CloudAccountRepo.getAccounts();
|
||||
const now = Math.floor(Date.now() / 1000);
|
||||
|
||||
for (const account of accounts) {
|
||||
try {
|
||||
// 1. Check/Refresh Token if needed (give it a 10 min buffer here for safety)
|
||||
let accessToken = account.token.access_token;
|
||||
if (account.token.expiry_timestamp < now + 600) {
|
||||
logger.info(`Monitor: Refreshing token for ${account.email}`);
|
||||
const newToken = await GoogleAPIService.refreshAccessToken(account.token.refresh_token);
|
||||
account.token.access_token = newToken.access_token;
|
||||
account.token.expires_in = newToken.expires_in;
|
||||
account.token.expiry_timestamp = now + newToken.expires_in;
|
||||
await CloudAccountRepo.updateToken(account.id, account.token);
|
||||
accessToken = newToken.access_token;
|
||||
}
|
||||
|
||||
// 2. Fetch Quota
|
||||
// We delay slightly between requests to act human/avoid spike
|
||||
await new Promise((r) => setTimeout(r, 1000));
|
||||
const quota = await GoogleAPIService.fetchQuota(accessToken);
|
||||
|
||||
// 3. Update DB
|
||||
await CloudAccountRepo.updateQuota(account.id, quota);
|
||||
} catch (error) {
|
||||
logger.error(`Monitor: Failed to update ${account.email}`, error);
|
||||
// Could mark status as 'error' or 'rate_limited' if 429
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Check for Auto-Switch
|
||||
await AutoSwitchService.checkAndSwitchIfNeeded();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,301 @@
|
||||
import { ConfigManager } from '../ipc/config/manager';
|
||||
import { ProxyAgent } from 'undici';
|
||||
|
||||
// --- Constants & Config ---
|
||||
|
||||
const CLIENT_ID = '1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com';
|
||||
const CLIENT_SECRET = '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';
|
||||
|
||||
declare const fetch: any;
|
||||
|
||||
const URLS = {
|
||||
TOKEN: 'https://oauth2.googleapis.com/token',
|
||||
USER_INFO: 'https://www.googleapis.com/oauth2/v2/userinfo',
|
||||
AUTH: 'https://accounts.google.com/o/oauth2/v2/auth',
|
||||
QUOTA: 'https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels',
|
||||
LOAD_PROJECT: 'https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist',
|
||||
};
|
||||
|
||||
// Internal API masquerading
|
||||
const USER_AGENT = 'antigravity/1.11.3 Darwin/arm64';
|
||||
const REDIRECT_URI = 'http://localhost:8888/oauth-callback';
|
||||
|
||||
// --- Types ---
|
||||
|
||||
export interface TokenResponse {
|
||||
access_token: string;
|
||||
expires_in: number;
|
||||
token_type: string;
|
||||
refresh_token?: string;
|
||||
scope?: string;
|
||||
}
|
||||
|
||||
export interface UserInfo {
|
||||
id: string;
|
||||
email: string;
|
||||
verified_email: boolean;
|
||||
name: string;
|
||||
given_name: string;
|
||||
family_name: string;
|
||||
picture: string;
|
||||
}
|
||||
|
||||
export interface QuotaData {
|
||||
models: Record<
|
||||
string,
|
||||
{
|
||||
percentage: number;
|
||||
resetTime: string;
|
||||
}
|
||||
>;
|
||||
}
|
||||
|
||||
// Internal types for API parsing
|
||||
interface ModelInfoRaw {
|
||||
quotaInfo?: {
|
||||
remainingFraction?: number;
|
||||
resetTime?: string;
|
||||
};
|
||||
}
|
||||
|
||||
interface LoadProjectResponse {
|
||||
cloudaicompanionProject?: string;
|
||||
}
|
||||
|
||||
// --- Service Implementation ---
|
||||
|
||||
export class GoogleAPIService {
|
||||
private static getFetchOptions(): any {
|
||||
try {
|
||||
const config = ConfigManager.loadConfig();
|
||||
if (config.proxy?.upstream_proxy?.enabled && config.proxy.upstream_proxy.url) {
|
||||
return {
|
||||
dispatcher: new ProxyAgent(config.proxy.upstream_proxy.url),
|
||||
};
|
||||
}
|
||||
} catch (e) {
|
||||
// Fallback or log if config load fails (shouldn't happen usually)
|
||||
console.warn('[GoogleAPIService] Failed to load proxy config', e);
|
||||
}
|
||||
return {};
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates the OAuth2 authorization URL.
|
||||
*/
|
||||
static getAuthUrl(): string {
|
||||
const scopes = [
|
||||
'https://www.googleapis.com/auth/cloud-platform',
|
||||
'https://www.googleapis.com/auth/userinfo.email',
|
||||
'https://www.googleapis.com/auth/userinfo.profile',
|
||||
'https://www.googleapis.com/auth/cclog',
|
||||
'https://www.googleapis.com/auth/experimentsandconfigs',
|
||||
].join(' ');
|
||||
|
||||
const params = new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
redirect_uri: REDIRECT_URI,
|
||||
response_type: 'code',
|
||||
scope: scopes,
|
||||
access_type: 'offline',
|
||||
prompt: 'consent',
|
||||
include_granted_scopes: 'true',
|
||||
});
|
||||
|
||||
return `${URLS.AUTH}?${params.toString()}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Exchanges an authorization code for tokens.
|
||||
*/
|
||||
static async exchangeCode(code: string): Promise<TokenResponse> {
|
||||
const params = new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
code: code,
|
||||
redirect_uri: REDIRECT_URI,
|
||||
grant_type: 'authorization_code',
|
||||
});
|
||||
|
||||
const response = await fetch(URLS.TOKEN, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
|
||||
body: params,
|
||||
...this.getFetchOptions(),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text();
|
||||
throw new Error(`Token exchange failed: ${text}`);
|
||||
}
|
||||
|
||||
return response.json() as Promise<TokenResponse>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Refreshes an access token using a refresh token.
|
||||
*/
|
||||
static async refreshAccessToken(refreshToken: string): Promise<TokenResponse> {
|
||||
const params = new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
refresh_token: refreshToken,
|
||||
grant_type: 'refresh_token',
|
||||
});
|
||||
|
||||
const response = await fetch(URLS.TOKEN, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
|
||||
body: params,
|
||||
...this.getFetchOptions(),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text();
|
||||
throw new Error(`Token refresh failed: ${text}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as TokenResponse;
|
||||
|
||||
return data;
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetches user profile information.
|
||||
*/
|
||||
static async getUserInfo(accessToken: string): Promise<UserInfo> {
|
||||
const response = await fetch(URLS.USER_INFO, {
|
||||
headers: { Authorization: `Bearer ${accessToken}` },
|
||||
...this.getFetchOptions(),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text();
|
||||
throw new Error(`Failed to fetch user info: ${text}`);
|
||||
}
|
||||
|
||||
return response.json() as Promise<UserInfo>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Core logic: Fetches the internal project ID needed for quota checks.
|
||||
*/
|
||||
public static async fetchProjectId(accessToken: string): Promise<string | null> {
|
||||
const body = {
|
||||
metadata: { ideType: 'ANTIGRAVITY' },
|
||||
};
|
||||
|
||||
// Simple retry logic (2 attempts)
|
||||
for (let i = 0; i < 2; i++) {
|
||||
try {
|
||||
const response = await fetch(URLS.LOAD_PROJECT, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
'User-Agent': USER_AGENT,
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
...this.getFetchOptions(),
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const data = (await response.json()) as LoadProjectResponse;
|
||||
if (data.cloudaicompanionProject) {
|
||||
return data.cloudaicompanionProject;
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn(`[GoogleAPIService] Failed to fetch project ID (Attempt ${i + 1}):`, e);
|
||||
await new Promise((r) => setTimeout(r, 500)); // Sleep 500ms
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Core logic: Fetches detailed model quota information.
|
||||
*/
|
||||
static async fetchQuota(accessToken: string): Promise<QuotaData> {
|
||||
// 1. Get Project ID (Critical)
|
||||
const projectId = await this.fetchProjectId(accessToken);
|
||||
|
||||
// 2. Build Payload
|
||||
const payload: Record<string, any> = {};
|
||||
if (projectId) {
|
||||
payload['project'] = projectId;
|
||||
}
|
||||
|
||||
// 3. Send Request with Retries
|
||||
const maxRetries = 3;
|
||||
let lastError: Error | null = null;
|
||||
const fetchOptions = this.getFetchOptions(); // Reuse options
|
||||
|
||||
for (let attempt = 1; attempt <= maxRetries; attempt++) {
|
||||
try {
|
||||
const response = await fetch(URLS.QUOTA, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
'User-Agent': USER_AGENT,
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify(payload),
|
||||
...fetchOptions,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text();
|
||||
const status = response.status;
|
||||
|
||||
if (status === 403) {
|
||||
throw new Error('FORBIDDEN');
|
||||
}
|
||||
if (status === 401) {
|
||||
throw new Error('UNAUTHORIZED');
|
||||
}
|
||||
|
||||
const errorMsg = `HTTP ${status} - ${text}`;
|
||||
console.warn(
|
||||
`[GoogleAPIService] API Error: ${errorMsg} (Attempt ${attempt}/${maxRetries})`,
|
||||
);
|
||||
|
||||
if (attempt < maxRetries) {
|
||||
await new Promise((r) => setTimeout(r, 1000));
|
||||
continue;
|
||||
} else {
|
||||
throw new Error(errorMsg);
|
||||
}
|
||||
}
|
||||
|
||||
const data = (await response.json()) as { models: Record<string, ModelInfoRaw> };
|
||||
const result: QuotaData = { models: {} };
|
||||
|
||||
// Parse relevant models
|
||||
for (const [name, info] of Object.entries(data.models || {})) {
|
||||
if (info.quotaInfo) {
|
||||
const fraction = info.quotaInfo.remainingFraction ?? 0;
|
||||
const percentage = Math.floor(fraction * 100);
|
||||
const resetTime = info.quotaInfo.resetTime || '';
|
||||
|
||||
if (name.toLowerCase().includes('gemini') || name.toLowerCase().includes('claude')) {
|
||||
result.models[name] = { percentage, resetTime };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
} catch (e: any) {
|
||||
console.warn(
|
||||
`[GoogleAPIService] Request failed: ${e.message} (Attempt ${attempt}/${maxRetries})`,
|
||||
);
|
||||
lastError = e;
|
||||
if (attempt < maxRetries) {
|
||||
await new Promise((r) => setTimeout(r, 1000));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
throw lastError || new Error('Quota check failed');
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user