torbjorn/antigravity-claude-proxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

code quality improvements and refactoring

Browse Source
This commit is contained in:
Badri Narayanan S
2025-12-21 22:11:44 +05:30
parent f282b36d1e
commit 712da8f7f2
5 changed files with 81 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
src/account-manager.js
View File

@@ -4,7 +4,8 @@
* automatic failover, and smart cooldown for rate-limited accounts.
*/
import { existsSync, readFileSync, writeFileSync, mkdirSync } from 'fs';
import { readFile, writeFile, mkdir, access } from 'fs/promises';
import { constants as fsConstants } from 'fs';
import { dirname } from 'path';
import { execSync } from 'child_process';
import { homedir } from 'os';
@@ -71,33 +72,34 @@ export class AccountManager {
if (this.#initialized) return;
try {
if (existsSync(this.#configPath)) {
const configData = readFileSync(this.#configPath, 'utf-8');
const config = JSON.parse(configData);
// Check if config file exists using async access
await access(this.#configPath, fsConstants.F_OK);
const configData = await readFile(this.#configPath, 'utf-8');
const config = JSON.parse(configData);
this.#accounts = (config.accounts || []).map(acc => ({
...acc,
isRateLimited: acc.isRateLimited || false,
rateLimitResetTime: acc.rateLimitResetTime || null,
lastUsed: acc.lastUsed || null
}));
this.#accounts = (config.accounts || []).map(acc => ({
...acc,
isRateLimited: acc.isRateLimited || false,
rateLimitResetTime: acc.rateLimitResetTime || null,
lastUsed: acc.lastUsed || null
}));
this.#settings = config.settings || {};
this.#currentIndex = config.activeIndex || 0;
this.#settings = config.settings || {};
this.#currentIndex = config.activeIndex || 0;
// Clamp currentIndex to valid range
if (this.#currentIndex >= this.#accounts.length) {
this.#currentIndex = 0;
}
// Clamp currentIndex to valid range
if (this.#currentIndex >= this.#accounts.length) {
this.#currentIndex = 0;
}
console.log(`[AccountManager] Loaded ${this.#accounts.length} account(s) from config`);
} else {
console.log(`[AccountManager] Loaded ${this.#accounts.length} account(s) from config`);
} catch (error) {
if (error.code === 'ENOENT') {
// No config file - use single account from Antigravity database
console.log('[AccountManager] No config file found. Using Antigravity database (single account mode)');
await this.#loadDefaultAccount();
} else {
console.error('[AccountManager] Failed to load config:', error.message);
}
} catch (error) {
console.error('[AccountManager] Failed to load config:', error.message);
// Fall back to default account
await this.#loadDefaultAccount();
}
@@ -341,7 +343,7 @@ export class AccountManager {
if (account.isInvalid) {
account.isInvalid = false;
account.invalidReason = null;
this.saveToDisk();
await this.saveToDisk();
}
console.log(`[AccountManager] Refreshed OAuth token for: ${account.email}`);
} catch (error) {
@@ -454,15 +456,13 @@ export class AccountManager {
}
/**
* Save current state to disk
* Save current state to disk (async)
*/
saveToDisk() {
async saveToDisk() {
try {
// Ensure directory exists
const dir = dirname(this.#configPath);
if (!existsSync(dir)) {
mkdirSync(dir, { recursive: true });
}
await mkdir(dir, { recursive: true });
const config = {
accounts: this.#accounts.map(acc => ({
@@ -483,7 +483,7 @@ export class AccountManager {
activeIndex: this.#currentIndex
};
writeFileSync(this.#configPath, JSON.stringify(config, null, 2));
await writeFile(this.#configPath, JSON.stringify(config, null, 2));
} catch (error) {
console.error('[AccountManager] Failed to save config:', error.message);
}

14
src/cloudcode-client.js
View File

@@ -14,7 +14,9 @@ import {
ANTIGRAVITY_ENDPOINT_FALLBACKS,
ANTIGRAVITY_HEADERS,
AVAILABLE_MODELS,
MAX_RETRIES
MAX_RETRIES,
MAX_WAIT_BEFORE_ERROR_MS,
MIN_SIGNATURE_LENGTH
} from './constants.js';
import {
mapModelName,
@@ -250,7 +252,7 @@ export async function sendMessage(anthropicRequest, accountManager) {
const resetTime = new Date(Date.now() + waitMs).toISOString();
// If wait time is too long (> 2 minutes), throw error immediately
if (waitMs > 120000) {
if (waitMs > MAX_WAIT_BEFORE_ERROR_MS) {
throw new Error(
`RESOURCE_EXHAUSTED: Rate limited. Quota will reset after ${formatDuration(waitMs)}. Next available: ${resetTime}`
);
@@ -448,7 +450,9 @@ async function parseThinkingSSEResponse(response, originalModel) {
accumulatedText += part.text;
}
}
} catch (e) { /* skip parse errors */ }
} catch (e) {
console.log('[CloudCode] SSE parse warning:', e.message, 'Raw:', jsonText.slice(0, 100));
}
}
}
@@ -496,7 +500,7 @@ export async function* sendMessageStream(anthropicRequest, accountManager) {
const resetTime = new Date(Date.now() + waitMs).toISOString();
// If wait time is too long (> 2 minutes), throw error immediately
if (waitMs > 120000) {
if (waitMs > MAX_WAIT_BEFORE_ERROR_MS) {
throw new Error(
`RESOURCE_EXHAUSTED: Rate limited. Quota will reset after ${formatDuration(waitMs)}. Next available: ${resetTime}`
);
@@ -692,7 +696,7 @@ async function* streamSSEResponse(response, originalModel) {
};
}
if (signature && signature.length >= 50) {
if (signature && signature.length >= MIN_SIGNATURE_LENGTH) {
currentThinkingSignature = signature;
}

14
src/constants.js
View File

@@ -81,6 +81,14 @@ export const ACCOUNT_CONFIG_PATH = join(
export const DEFAULT_COOLDOWN_MS = 60 * 1000; // 1 minute default cooldown
export const MAX_RETRIES = 5; // Max retry attempts across accounts
// Rate limit wait thresholds
export const MAX_WAIT_BEFORE_ERROR_MS = 120000; // 2 minutes - throw error if wait exceeds this
// Thinking model constants
export const DEFAULT_THINKING_BUDGET = 16000; // Default thinking budget tokens
export const CLAUDE_THINKING_MAX_OUTPUT_TOKENS = 64000; // Max output tokens for thinking models
export const MIN_SIGNATURE_LENGTH = 50; // Minimum valid thinking signature length
export default {
ANTIGRAVITY_ENDPOINT_FALLBACKS,
ANTIGRAVITY_HEADERS,
@@ -93,5 +101,9 @@ export default {
DEFAULT_PORT,
ACCOUNT_CONFIG_PATH,
DEFAULT_COOLDOWN_MS,
MAX_RETRIES
MAX_RETRIES,
MAX_WAIT_BEFORE_ERROR_MS,
DEFAULT_THINKING_BUDGET,
CLAUDE_THINKING_MAX_OUTPUT_TOKENS,
MIN_SIGNATURE_LENGTH
};

22
src/format-converter.js
View File

@@ -8,12 +8,12 @@
*/
import crypto from 'crypto';
import { MODEL_MAPPINGS } from './constants.js';
// Default thinking budget (16K tokens)
const DEFAULT_THINKING_BUDGET = 16000;
// Claude thinking models need larger max output tokens
const CLAUDE_THINKING_MAX_OUTPUT_TOKENS = 64000;
import {
MODEL_MAPPINGS,
DEFAULT_THINKING_BUDGET,
CLAUDE_THINKING_MAX_OUTPUT_TOKENS,
MIN_SIGNATURE_LENGTH
} from './constants.js';
/**
* Map Anthropic model name to Antigravity model name
@@ -33,11 +33,11 @@ function isThinkingPart(part) {
}
/**
* Check if a thinking part has a valid signature (>= 50 chars)
* Check if a thinking part has a valid signature (>= MIN_SIGNATURE_LENGTH chars)
*/
function hasValidSignature(part) {
const signature = part.thought === true ? part.thoughtSignature : part.signature;
return typeof signature === 'string' && signature.length >= 50;
return typeof signature === 'string' && signature.length >= MIN_SIGNATURE_LENGTH;
}
/**
@@ -187,8 +187,8 @@ export function restoreThinkingSignatures(content) {
continue;
}
// Keep blocks with valid signatures (>= 50 chars), sanitized
if (block.signature && block.signature.length >= 50) {
// Keep blocks with valid signatures (>= MIN_SIGNATURE_LENGTH chars), sanitized
if (block.signature && block.signature.length >= MIN_SIGNATURE_LENGTH) {
filtered.push(sanitizeAnthropicThinkingBlock(block));
}
// Unsigned thinking blocks are dropped
@@ -360,7 +360,7 @@ function convertContentToParts(content, isClaudeModel = false) {
parts.push({ functionResponse });
} else if (block.type === 'thinking') {
// Handle thinking blocks - only those with valid signatures
if (block.signature && block.signature.length >= 50) {
if (block.signature && block.signature.length >= MIN_SIGNATURE_LENGTH) {
// Convert to Gemini format with signature
parts.push({
text: block.thinking,

31
src/server.js
View File

@@ -19,23 +19,32 @@ const accountManager = new AccountManager();
// Track initialization status
let isInitialized = false;
let initError = null;
let initPromise = null;
/**
* Ensure account manager is initialized
* Ensure account manager is initialized (with race condition protection)
*/
async function ensureInitialized() {
if (isInitialized) return;
try {
await accountManager.initialize();
isInitialized = true;
const status = accountManager.getStatus();
console.log(`[Server] Account pool initialized: ${status.summary}`);
} catch (error) {
initError = error;
console.error('[Server] Failed to initialize account manager:', error.message);
throw error;
}
// If initialization is already in progress, wait for it
if (initPromise) return initPromise;
initPromise = (async () => {
try {
await accountManager.initialize();
isInitialized = true;
const status = accountManager.getStatus();
console.log(`[Server] Account pool initialized: ${status.summary}`);
} catch (error) {
initError = error;
initPromise = null; // Allow retry on failure
console.error('[Server] Failed to initialize account manager:', error.message);
throw error;
}
})();
return initPromise;
}
// Middleware