luulam/9router
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Feat : Auto restart after crash

Browse Source
This commit is contained in:
decolua
2026-03-14 09:37:29 +07:00
parent 549223c8cf
commit adae2605bf
31 changed files with 340 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/app/api/usage/stream/route.js
View File

@@ -39,7 +39,9 @@ export async function GET() {
controller.enqueue(encoder.encode(`data: ${JSON.stringify(stats)}\n\n`));
} catch {
state.closed = true;
statsEmitter.off("update", state.send);
statsEmitter.off("pending", state.sendPending);
clearInterval(state.keepalive);
}
};

28
src/lib/requestDetailsDb.js
View File

@@ -246,23 +246,25 @@ export async function getRequestDetailById(id) {
return db.data.records.find(r => r.id === id) || null;
}
// Graceful shutdown
let shutdownHandlerRegistered = false;
// Graceful shutdown — use named handler so we can remove it on re-registration
const _shutdownHandler = async () => {
if (flushTimer) { clearTimeout(flushTimer); flushTimer = null; }
if (writeBuffer.length > 0) await flushToDatabase();
};
function ensureShutdownHandler() {
if (shutdownHandlerRegistered || isCloud) return;
if (isCloud) return;
const handler = async () => {
if (flushTimer) { clearTimeout(flushTimer); flushTimer = null; }
if (writeBuffer.length > 0) await flushToDatabase();
};
// Remove any previously registered listeners from this module (hot-reload safety)
process.off("beforeExit", _shutdownHandler);
process.off("SIGINT", _shutdownHandler);
process.off("SIGTERM", _shutdownHandler);
process.off("exit", _shutdownHandler);
process.on("beforeExit", handler);
process.on("SIGINT", handler);
process.on("SIGTERM", handler);
process.on("exit", handler);
shutdownHandlerRegistered = true;
process.on("beforeExit", _shutdownHandler);
process.on("SIGINT", _shutdownHandler);
process.on("SIGTERM", _shutdownHandler);
process.on("exit", _shutdownHandler);
}
ensureShutdownHandler();

11
src/lib/tunnel/cloudflared.js
View File

@@ -141,8 +141,10 @@ export async function spawnCloudflared(tunnelToken) {
const handleLog = (data) => {
const msg = data.toString();
if (msg.includes("Registered tunnel connection")) {
connectionCount++;
// Count exact occurrences in this chunk (each chunk may contain multiple lines)
const matches = msg.match(/Registered tunnel connection/g);
if (matches) {
connectionCount += matches.length;
if (connectionCount >= 4 && !resolved) {
resolved = true;
clearTimeout(timeout);
@@ -165,6 +167,7 @@ export async function spawnCloudflared(tunnelToken) {
child.on("exit", (code) => {
cloudflaredProcess = null;
clearPid();
const wasConnected = resolved; // true = already connected successfully
if (!resolved) {
resolved = true;
clearTimeout(timeout);
@@ -173,8 +176,8 @@ export async function spawnCloudflared(tunnelToken) {
return;
}
}
// Notify reconnect handler if tunnel died after successful connection
if (unexpectedExitHandler) {
// Only notify on unexpected exit AFTER successful connection
if (wasConnected && unexpectedExitHandler) {
unexpectedExitHandler();
}
});

11
src/lib/tunnel/tunnelManager.js
View File

@@ -8,7 +8,8 @@ const MACHINE_ID_SALT = "9router-tunnel-salt";
const API_KEY_SECRET = "9router-tunnel-api-key-secret";
const SHORT_ID_LENGTH = 6;
const SHORT_ID_CHARS = "abcdefghijklmnpqrstuvwxyz23456789";
const RECONNECT_DELAYS_MS = [5000, 15000, 30000];
const RECONNECT_DELAYS_MS = [5000, 10000, 20000, 30000, 60000];
const MAX_RECONNECT_ATTEMPTS = RECONNECT_DELAYS_MS.length;
let isReconnecting = false;
@@ -83,8 +84,10 @@ export async function enableTunnel() {
await updateSettings({ tunnelEnabled: true, tunnelUrl: hostname });
// Register exit handler for auto-reconnect on unexpected crash/sleep-wake
setUnexpectedExitHandler(() => scheduleReconnect(0));
// Re-register exit handler each time tunnel starts (handles reconnect scenario too)
setUnexpectedExitHandler(() => {
if (!isReconnecting) scheduleReconnect(0);
});
return { success: true, tunnelUrl: hostname, shortId };
}
@@ -112,7 +115,7 @@ async function scheduleReconnect(attempt) {
console.log(`[Tunnel] Reconnect attempt ${attempt + 1} failed:`, err.message);
isReconnecting = false;
const nextAttempt = attempt + 1;
if (nextAttempt < RECONNECT_DELAYS_MS.length) {
if (nextAttempt < MAX_RECONNECT_ATTEMPTS) {
scheduleReconnect(nextAttempt);
} else {
console.log("[Tunnel] All reconnect attempts exhausted");

7
src/lib/usageDb.js
View File

@@ -245,8 +245,11 @@ export async function saveRequestUsage(entry) {
entry.cost = entryCost;
db.data.history.push(entry);
// Optional: Limit history size if needed in future
// if (db.data.history.length > 10000) db.data.history.shift();
// Cap history to prevent unbounded memory/disk growth
const MAX_HISTORY = 10000;
if (db.data.history.length > MAX_HISTORY) {
db.data.history.splice(0, db.data.history.length - MAX_HISTORY);
}
await db.write();
statsEmitter.emit("update");

59
src/mitm/manager.js
View File

@@ -16,6 +16,14 @@ const MITM_PORT = 443;
const MITM_WIN_NODE_PORT = 8443;
const PID_FILE = path.join(MITM_DIR, ".mitm.pid");
const MITM_MAX_RESTARTS = 5;
const MITM_RESTART_DELAYS_MS = [5000, 10000, 20000, 30000, 60000];
const MITM_RESTART_RESET_MS = 60000;
let mitmRestartCount = 0;
let mitmLastStartTime = 0;
let mitmIsRestarting = false;
function resolveServerPath() {
if (process.env.MITM_SERVER_PATH) return process.env.MITM_SERVER_PATH;
const sibling = path.join(__dirname, "server.js");
@@ -273,6 +281,50 @@ async function getMitmStatus() {
return { running, pid, certExists, dnsStatus };
}
async function scheduleMitmRestart(apiKey) {
if (mitmIsRestarting) return;
const aliveMs = Date.now() - mitmLastStartTime;
if (aliveMs >= MITM_RESTART_RESET_MS) mitmRestartCount = 0;
if (mitmRestartCount >= MITM_MAX_RESTARTS) {
console.error("[MITM] Max restart attempts reached. Giving up.");
return;
}
const attempt = mitmRestartCount;
const delay = MITM_RESTART_DELAYS_MS[Math.min(attempt, MITM_RESTART_DELAYS_MS.length - 1)];
mitmRestartCount++;
mitmIsRestarting = true;
console.log(`[MITM] Restarting in ${delay / 1000}s... (${mitmRestartCount}/${MITM_MAX_RESTARTS})`);
await new Promise((r) => setTimeout(r, delay));
try {
const settings = _getSettings ? await _getSettings() : null;
if (settings && !settings.mitmEnabled) {
console.log("[MITM] MITM disabled, skipping restart");
mitmIsRestarting = false;
return;
}
const password = getCachedPassword() || await loadEncryptedPassword();
if (!password && !IS_WIN) {
console.error("[MITM] No cached password, cannot auto-restart");
mitmIsRestarting = false;
return;
}
await startServer(apiKey, password);
console.log("[MITM] Restarted successfully");
mitmRestartCount = 0;
mitmIsRestarting = false;
} catch (err) {
console.error(`[MITM] Restart attempt ${mitmRestartCount}/${MITM_MAX_RESTARTS} failed:`, err.message);
mitmIsRestarting = false;
// Schedule next retry
scheduleMitmRestart(apiKey);
}
}
/**
* Start MITM server only (cert + server, no DNS)
*/
@@ -378,6 +430,7 @@ async function startServer(apiKey, sudoPassword) {
if (!IS_WIN && serverProcess) {
serverPid = serverProcess.pid;
fs.writeFileSync(PID_FILE, String(serverPid));
mitmLastStartTime = Date.now();
}
let startError = null;
@@ -397,6 +450,8 @@ async function startServer(apiKey, sudoPassword) {
serverProcess = null;
serverPid = null;
try { fs.unlinkSync(PID_FILE); } catch { /* ignore */ }
// Auto-restart on unexpected exit
if (code !== 0 && !mitmIsRestarting) scheduleMitmRestart(apiKey);
});
}
@@ -425,6 +480,9 @@ async function startServer(apiKey, sudoPassword) {
* Stop MITM server — removes ALL tool DNS entries first, then kills server
*/
async function stopServer(sudoPassword) {
// Prevent auto-restart from triggering on intentional stop
mitmIsRestarting = true;
mitmRestartCount = 0;
console.log("[MITM] Stopping server...");
// Kill server process
@@ -476,6 +534,7 @@ async function stopServer(sudoPassword) {
try { fs.unlinkSync(PID_FILE); } catch { /* ignore */ }
await saveMitmSettings(false, null);
mitmIsRestarting = false;
return { running: false, pid: null };
}

2
src/mitm/server.js
View File

@@ -85,7 +85,7 @@ const ANTIGRAVITY_URL_PATTERNS = [":generateContent", ":streamGenerateContent"];
// Copilot: OpenAI-compatible + Anthropic endpoints
const COPILOT_URL_PATTERNS = ["/chat/completions", "/v1/messages", "/responses"];
const LOG_DIR = path.join(__dirname, "../../logs/mitm");
const LOG_DIR = path.join(DATA_DIR, "logs", "mitm");
if (ENABLE_FILE_LOG && !fs.existsSync(LOG_DIR)) fs.mkdirSync(LOG_DIR, { recursive: true });
function saveRequestLog(url, bodyBuffer) {

10
src/sse/handlers/chat.js
View File

@@ -111,7 +111,7 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re
return handleComboChat({
body,
models: comboModels,
handleSingleModel: (b, m) => handleSingleModelChat(b, m, clientRawRequest, request, apiKey, forceSourceFormat),
handleSingleModel: (b, m) => handleSingleModelChat(b, m, clientRawRequest, request, apiKey),
log
});
}
@@ -132,12 +132,12 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re
const userAgent = request?.headers?.get("user-agent") || "";
// Try with available accounts (fallback on errors)
let excludeConnectionId = null;
const excludeConnectionIds = new Set();
let lastError = null;
let lastStatus = null;
while (true) {
const credentials = await getProviderCredentials(provider, excludeConnectionId, model);
const credentials = await getProviderCredentials(provider, excludeConnectionIds, model);
// All accounts unavailable
if (!credentials || credentials.allRateLimited) {
@@ -147,7 +147,7 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re
log.warn("CHAT", `[${provider}/${model}] ${errorMsg} (${credentials.retryAfterHuman})`);
return unavailableResponse(status, `[${provider}/${model}] ${errorMsg}`, credentials.retryAfter, credentials.retryAfterHuman);
}
if (!excludeConnectionId) {
if (excludeConnectionIds.size === 0) {
log.error("AUTH", `No credentials for provider: ${provider}`);
return errorResponse(HTTP_STATUS.BAD_REQUEST, `No credentials for provider: ${provider}`);
}
@@ -204,7 +204,7 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re
if (shouldFallback) {
log.warn("AUTH", `Account ${credentials.connectionName} unavailable (${result.status}), trying fallback`);
excludeConnectionId = credentials.connectionId;
excludeConnectionIds.add(credentials.connectionId);
lastError = result.error;
lastStatus = result.status;
continue;

8
src/sse/handlers/embeddings.js
View File

@@ -80,12 +80,12 @@ export async function handleEmbeddings(request) {
}
// Credential + fallback loop (mirrors handleChat)
let excludeConnectionId = null;
const excludeConnectionIds = new Set();
let lastError = null;
let lastStatus = null;
while (true) {
const credentials = await getProviderCredentials(provider, excludeConnectionId, model);
const credentials = await getProviderCredentials(provider, excludeConnectionIds, model);
// All accounts unavailable
if (!credentials || credentials.allRateLimited) {
@@ -95,7 +95,7 @@ export async function handleEmbeddings(request) {
log.warn("EMBEDDINGS", `[${provider}/${model}] ${errorMsg} (${credentials.retryAfterHuman})`);
return unavailableResponse(status, `[${provider}/${model}] ${errorMsg}`, credentials.retryAfter, credentials.retryAfterHuman);
}
if (!excludeConnectionId) {
if (excludeConnectionIds.size === 0) {
log.error("AUTH", `No credentials for provider: ${provider}`);
return errorResponse(HTTP_STATUS.BAD_REQUEST, `No credentials for provider: ${provider}`);
}
@@ -131,7 +131,7 @@ export async function handleEmbeddings(request) {
if (shouldFallback) {
log.warn("AUTH", `Account ${credentials.connectionName} unavailable (${result.status}), trying fallback`);
excludeConnectionId = credentials.connectionId;
excludeConnectionIds.add(credentials.connectionId);
lastError = result.error;
lastStatus = result.status;
continue;

14
src/sse/services/auth.js
View File

@@ -11,10 +11,14 @@ let selectionMutex = Promise.resolve();
* Get provider credentials from localDb
* Filters out unavailable accounts and returns the selected account based on strategy
* @param {string} provider - Provider name
* @param {string|null} excludeConnectionId - Connection ID to exclude (for retry with next account)
* @param {Set<string>|string|null} excludeConnectionIds - Connection ID(s) to exclude (for retry with next account)
* @param {string|null} model - Model name for per-model rate limit filtering
*/
export async function getProviderCredentials(provider, excludeConnectionId = null, model = null) {
export async function getProviderCredentials(provider, excludeConnectionIds = null, model = null) {
// Normalize to Set for consistent handling
const excludeSet = excludeConnectionIds instanceof Set
? excludeConnectionIds
: (excludeConnectionIds ? new Set([excludeConnectionIds]) : new Set());
// Acquire mutex to prevent race conditions
const currentMutex = selectionMutex;
let resolveMutex;
@@ -27,7 +31,7 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul
const providerId = resolveProviderId(provider);
const connections = await getProviderConnections({ provider: providerId, isActive: true });
log.debug("AUTH", `${provider} | total connections: ${connections.length}, excludeId: ${excludeConnectionId || "none"}, model: ${model || "any"}`);
log.debug("AUTH", `${provider} | total connections: ${connections.length}, excludeIds: ${excludeSet.size > 0 ? [...excludeSet].join(",") : "none"}, model: ${model || "any"}`);
if (connections.length === 0) {
log.warn("AUTH", `No credentials for ${provider}`);
@@ -36,14 +40,14 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul
// Filter out model-locked and excluded connections
const availableConnections = connections.filter(c => {
if (excludeConnectionId && c.id === excludeConnectionId) return false;
if (excludeSet.has(c.id)) return false;
if (isModelLockActive(c, model)) return false;
return true;
});
log.debug("AUTH", `${provider} | available: ${availableConnections.length}/${connections.length}`);
connections.forEach(c => {
const excluded = excludeConnectionId && c.id === excludeConnectionId;
const excluded = excludeSet.has(c.id);
const locked = isModelLockActive(c, model);
if (excluded || locked) {
const lockUntil = getEarliestModelLockUntil(c);