OmniRoute/tests/unit/fixes-p1.test.ts

import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
import net from "node:net";

const isWindows = process.platform === "win32";
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-fixes-"));
process.env.DATA_DIR = TEST_DATA_DIR;

const core = await import("../../src/lib/db/core.ts");
const backupDb = await import("../../src/lib/db/backup.ts");
const providersDb = await import("../../src/lib/db/providers.ts");
const combosDb = await import("../../src/lib/db/combos.ts");
const settingsDb = await import("../../src/lib/db/settings.ts");
const localDb = await import("../../src/lib/localDb.ts");
const tokenRefresh = await import("../../open-sse/services/tokenRefresh.ts");
const proxyFetch = await import("../../open-sse/utils/proxyFetch.ts");
const proxyDispatcher = await import("../../open-sse/utils/proxyDispatcher.ts");
const proxySettingsRoute = await import("../../src/app/api/settings/proxy/route.ts");
const proxyTestRoute = await import("../../src/app/api/settings/proxy/test/route.ts");
const shutdownRoute = await import("../../src/app/api/shutdown/route.ts");
const restartRoute = await import("../../src/app/api/restart/route.ts");

async function withEnv(name, value, fn) {
  const previous = process.env[name];
  if (value === undefined) {
    delete process.env[name];
  } else {
    process.env[name] = value;
  }

  try {
    return await fn();
  } finally {
    if (previous === undefined) {
      delete process.env[name];
    } else {
      process.env[name] = previous;
    }
  }
}

async function resetStorage() {
  core.resetDbInstance();
  for (let attempt = 0; attempt < 10; attempt++) {
    try {
      if (fs.existsSync(TEST_DATA_DIR)) {
        fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
      }
      break;
    } catch (err) {
      if ((err?.code === "EBUSY" || err?.code === "EPERM") && attempt < 9) {
        await new Promise((r) => setTimeout(r, 100 * (attempt + 1)));
      } else {
        throw err;
      }
    }
  }
  fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
}

test.after(async () => {
  core.resetDbInstance();
  fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
});

test("token refresh dedupe key avoids collision for same-prefix tokens", async () => {
  const originalFetch = globalThis.fetch;
  const requests = [];

  globalThis.fetch = async (_url, options = {}) => {
    const refreshToken = options.body?.get?.("refresh_token") || "unknown";
    requests.push(refreshToken);
    return new Response(
      JSON.stringify({
        access_token: `access-${refreshToken}`,
        refresh_token: `refresh-${refreshToken}`,
        expires_in: 3600,
      }),
      {
        status: 200,
        headers: { "content-type": "application/json" },
      }
    );
  };

  try {
    const prefix = "abcdefghijklmnop";
    const tokenA = `${prefix}-account-a`;
    const tokenB = `${prefix}-account-b`;

    const [resultA, resultB] = await Promise.all([
      tokenRefresh.getAccessToken("codex", { refreshToken: tokenA }, null),
      tokenRefresh.getAccessToken("codex", { refreshToken: tokenB }, null),
    ]);

    assert.equal(requests.length, 2);
    assert.equal(resultA.accessToken, `access-${tokenA}`);
    assert.equal(resultB.accessToken, `access-${tokenB}`);
    assert.notEqual(resultA.accessToken, resultB.accessToken);
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test(
  "restoreDbBackup clears stale sqlite sidecars before reopen",
  { skip: isWindows },
  async () => {
    await resetStorage();

    const db = core.getDbInstance();
    const now = new Date().toISOString();
    db.prepare(
      "INSERT INTO provider_connections (id, provider, auth_type, name, is_active, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
    ).run("restore-test-conn", "openai", "apikey", "restore-test", 1, now, now);

    const backupId = "db_2000-01-01T00-00-00-000Z_manual.sqlite";
    const backupPath = path.join(core.DB_BACKUPS_DIR, backupId);
    fs.mkdirSync(core.DB_BACKUPS_DIR, { recursive: true });
    await db.backup(backupPath);

    core.resetDbInstance();
    fs.writeFileSync(`${core.SQLITE_FILE}-wal`, "STALE-WAL-MARKER");
    fs.writeFileSync(`${core.SQLITE_FILE}-shm`, "STALE-SHM-MARKER");
    fs.writeFileSync(`${core.SQLITE_FILE}-journal`, "STALE-JOURNAL-MARKER");

    await backupDb.restoreDbBackup(backupId);

    for (const suffix of ["-wal", "-shm", "-journal"]) {
      const sidecarPath = `${core.SQLITE_FILE}${suffix}`;
      if (!fs.existsSync(sidecarPath)) continue;
      const text = fs.readFileSync(sidecarPath, "utf8");
      assert.equal(text.includes("STALE-"), false, `sidecar ${suffix} still contains stale marker`);
    }

    const reopenedDb = core.getDbInstance();
    const row = reopenedDb
      .prepare("SELECT COUNT(*) AS cnt FROM provider_connections WHERE id = ?")
      .get("restore-test-conn");
    assert.equal(row.cnt, 1);
  }
);

test("closeDbInstance checkpoints WAL changes into the primary SQLite file", async () => {
  await resetStorage();

  const db = core.getDbInstance();
  const now = new Date().toISOString();
  db.prepare(
    "INSERT INTO provider_connections (id, provider, auth_type, name, is_active, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
  ).run("checkpoint-test-conn", "openai", "apikey", "checkpoint-test", 1, now, now);

  core.closeDbInstance();

  const snapshotPath = path.join(TEST_DATA_DIR, "storage-snapshot.sqlite");
  fs.copyFileSync(core.SQLITE_FILE, snapshotPath);

  const Database = (await import("better-sqlite3")).default;
  const snapshotDb = new Database(snapshotPath, { readonly: true });
  try {
    const row = snapshotDb
      .prepare("SELECT name FROM provider_connections WHERE id = ?")
      .get("checkpoint-test-conn");
    assert.equal(row?.name, "checkpoint-test");
  } finally {
    snapshotDb.close();
  }
});

test("shutdown route uses SIGTERM for graceful shutdown", async () => {
  const originalKill = process.kill;
  const originalSetTimeout = globalThis.setTimeout;
  const calls = [];

  process.kill = (pid, signal) => {
    calls.push({ pid, signal });
    return true;
  };
  globalThis.setTimeout = (callback) => {
    callback();
    return 0;
  };

  try {
    const response = await shutdownRoute.POST();
    assert.equal(response.status, 200);
    assert.deepEqual(calls, [{ pid: process.pid, signal: "SIGTERM" }]);
  } finally {
    process.kill = originalKill;
    globalThis.setTimeout = originalSetTimeout;
  }
});

test("restart route uses SIGTERM for graceful restart", async () => {
  const originalKill = process.kill;
  const originalSetTimeout = globalThis.setTimeout;
  const calls = [];

  process.kill = (pid, signal) => {
    calls.push({ pid, signal });
    return true;
  };
  globalThis.setTimeout = (callback) => {
    callback();
    return 0;
  };

  try {
    const response = await restartRoute.POST();
    assert.equal(response.status, 200);
    assert.deepEqual(calls, [{ pid: process.pid, signal: "SIGTERM" }]);
  } finally {
    process.kill = originalKill;
    globalThis.setTimeout = originalSetTimeout;
  }
});

test("unlinkFileWithRetry retries EBUSY/EPERM and eventually succeeds", async () => {
  const target = path.join(TEST_DATA_DIR, "retry-target.tmp");
  fs.writeFileSync(target, "retry-me");

  const originalExistsSync = fs.existsSync;
  const originalUnlinkSync = fs.unlinkSync;
  const seenCodes = [];
  let attempts = 0;

  fs.existsSync = (filePath) => {
    if (filePath === target) return attempts < 3 || originalExistsSync(filePath);
    return originalExistsSync(filePath);
  };

  fs.unlinkSync = (filePath) => {
    if (filePath === target) {
      attempts++;
      if (attempts === 1) {
        const err = new Error("busy");
        err.code = "EBUSY";
        seenCodes.push(err.code);
        throw err;
      }
      if (attempts === 2) {
        const err = new Error("perm");
        err.code = "EPERM";
        seenCodes.push(err.code);
        throw err;
      }
    }
    return originalUnlinkSync(filePath);
  };

  try {
    await backupDb.unlinkFileWithRetry(target, { maxAttempts: 5, baseDelayMs: 1 });
    assert.equal(attempts, 3);
    assert.deepEqual(seenCodes, ["EBUSY", "EPERM"]);
    assert.equal(fs.existsSync(target), false);
  } finally {
    fs.existsSync = originalExistsSync;
    fs.unlinkSync = originalUnlinkSync;
    if (originalExistsSync(target)) originalUnlinkSync(target);
  }
});

test("provider connection persists rateLimitProtection across reopen", async () => {
  await resetStorage();

  const created = await providersDb.createProviderConnection({
    provider: "openai",
    authType: "apikey",
    name: "rl-test",
    apiKey: "sk-test",
  });

  await providersDb.updateProviderConnection(created.id, { rateLimitProtection: true });

  const firstRead = await providersDb.getProviderConnectionById(created.id);
  assert.equal(firstRead.rateLimitProtection, true);

  core.resetDbInstance();
  const secondRead = await providersDb.getProviderConnectionById(created.id);
  assert.equal(secondRead.rateLimitProtection, true);
});

test('provider connection migration adds "group" column for existing databases', async () => {
  await resetStorage();

  const sqlitePath = core.SQLITE_FILE;
  core.resetDbInstance();

  const Database = (await import("better-sqlite3")).default;
  const db = new Database(sqlitePath);
  db.exec(`
    CREATE TABLE provider_connections (
      id TEXT PRIMARY KEY,
      provider TEXT NOT NULL,
      auth_type TEXT,
      name TEXT,
      email TEXT,
      priority INTEGER DEFAULT 0,
      is_active INTEGER DEFAULT 1,
      access_token TEXT,
      refresh_token TEXT,
      expires_at TEXT,
      token_expires_at TEXT,
      scope TEXT,
      project_id TEXT,
      test_status TEXT,
      error_code TEXT,
      last_error TEXT,
      last_error_at TEXT,
      last_error_type TEXT,
      last_error_source TEXT,
      backoff_level INTEGER DEFAULT 0,
      rate_limited_until TEXT,
      health_check_interval INTEGER,
      last_health_check_at TEXT,
      last_tested TEXT,
      api_key TEXT,
      id_token TEXT,
      provider_specific_data TEXT,
      expires_in INTEGER,
      display_name TEXT,
      global_priority INTEGER,
      default_model TEXT,
      token_type TEXT,
      consecutive_use_count INTEGER DEFAULT 0,
      rate_limit_protection INTEGER DEFAULT 0,
      last_used_at TEXT,
      created_at TEXT NOT NULL,
      updated_at TEXT NOT NULL
    )
  `);
  db.close();

  const reopened = core.getDbInstance();
  const columns = reopened.prepare("PRAGMA table_info(provider_connections)").all();
  const names = new Set(columns.map((column) => column.name));
  assert.equal(names.has("group"), true);
});

test("resolveProxyForConnection applies combo proxy for object/string model entries", async () => {
  await resetStorage();

  const conn = await providersDb.createProviderConnection({
    provider: "claude",
    authType: "oauth",
    email: "combo-test@example.com",
    accessToken: "access",
    refreshToken: "refresh",
  });

  const combo = await combosDb.createCombo({
    id: "combo-proxy-test",
    name: "combo-proxy-test",
    strategy: "priority",
    models: ["openai/gpt-5", { model: "cc/claude-sonnet-4-5-20250929", weight: 100 }],
  });

  await settingsDb.setProxyForLevel("combo", combo.id, {
    type: "http",
    host: "127.0.0.1",
    port: "8080",
  });

  const resolved = await settingsDb.resolveProxyForConnection(conn.id);
  assert.equal(resolved.level, "combo");
  assert.equal(resolved.levelId, combo.id);
});

test("normalizeProxyUrl accepts socks5 only when explicitly allowed", () => {
  const socksUrl = "socks5://127.0.0.1:1080";

  const normalized = proxyDispatcher.normalizeProxyUrl(socksUrl, "test proxy", {
    allowSocks5: true,
  });
  assert.match(normalized, /^socks5:\/\/127\.0\.0\.1:1080\/?$/);

  assert.throws(
    () =>
      proxyDispatcher.normalizeProxyUrl(socksUrl, "test proxy", {
        allowSocks5: false,
      }),
    /SOCKS5 proxy is disabled/i
  );
});

test("createProxyDispatcher builds dispatchers for http, https, and socks5", async () => {
  await withEnv("ENABLE_SOCKS5_PROXY", "true", async () => {
    const httpDispatcher = proxyDispatcher.createProxyDispatcher("http://127.0.0.1:8080");
    const httpsDispatcher = proxyDispatcher.createProxyDispatcher("https://127.0.0.1:8443");
    const socksDispatcher = proxyDispatcher.createProxyDispatcher("socks5://127.0.0.1:1080");

    assert.equal(typeof httpDispatcher.dispatch, "function");
    assert.equal(typeof httpsDispatcher.dispatch, "function");
    assert.equal(typeof socksDispatcher.dispatch, "function");
  });
});

test("proxy fetch fails closed when context proxy is invalid", async () => {
  await assert.rejects(
    proxyFetch.runWithProxyContext({ type: "http", host: "127.0.0.1", port: "9" }, async () =>
      fetch("https://example.com")
    )
  );
});

test("proxy fetch rejects socks5 context when feature flag is disabled", async () => {
  await withEnv("ENABLE_SOCKS5_PROXY", "false", async () => {
    await assert.rejects(
      proxyFetch.runWithProxyContext(
        { type: "socks5", host: "127.0.0.1", port: "1080" },
        async () => fetch("https://example.com")
      ),
      /ENABLE_SOCKS5_PROXY/i
    );
  });
});

test("proxy fetch accepts socks5 context when feature flag is enabled", async () => {
  await withEnv("ENABLE_SOCKS5_PROXY", "true", async () => {
    const server = net.createServer();
    await new Promise((resolve, reject) => {
      server.once("error", reject);
      server.listen(0, "127.0.0.1", resolve);
    });

    const address = server.address();
    assert.ok(address && typeof address === "object");

    try {
      const result = await proxyFetch.runWithProxyContext(
        { type: "socks5", host: "127.0.0.1", port: String(address.port) },
        async () => "ok"
      );
      assert.equal(result, "ok");
    } finally {
      await new Promise((resolve, reject) => {
        server.close((err) => {
          if (err) reject(err);
          else resolve();
        });
      });
    }
  });
});

test("proxy settings route blocks socks5 with backend flag disabled", async () => {
  await resetStorage();

  await withEnv("ENABLE_SOCKS5_PROXY", "false", async () => {
    const request = new Request("http://localhost/api/settings/proxy", {
      method: "PUT",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        level: "global",
        proxy: {
          type: "socks5",
          host: "127.0.0.1",
          port: "1080",
        },
      }),
    });

    const response = await proxySettingsRoute.PUT(request);
    const payload = await response.json();
    assert.equal(response.status, 400);
    assert.match(payload.error.message, /SOCKS5 proxy is disabled/i);
  });
});

test("proxy settings route accepts socks5 with backend flag enabled", async () => {
  await resetStorage();

  await withEnv("ENABLE_SOCKS5_PROXY", "true", async () => {
    const request = new Request("http://localhost/api/settings/proxy", {
      method: "PUT",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        level: "global",
        proxy: {
          type: "SOCKS5",
          host: "127.0.0.1",
          port: "1080",
        },
      }),
    });

    const response = await proxySettingsRoute.PUT(request);
    const payload = await response.json();
    assert.equal(response.status, 200);
    assert.equal(payload.global.type, "socks5");
  });
});

test("proxy test route rejects socks5 when backend flag is disabled", async () => {
  await withEnv("ENABLE_SOCKS5_PROXY", "false", async () => {
    const request = new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        proxy: {
          type: "socks5",
          host: "127.0.0.1",
          port: "1080",
        },
      }),
    });

    const response = await proxyTestRoute.POST(request);
    const payload = await response.json();

    assert.equal(response.status, 400);
    assert.match(payload.error.message, /SOCKS5 proxy is disabled/i);
  });
});

test("proxy test route runs socks5 test when backend flag is enabled", async () => {
  await withEnv("ENABLE_SOCKS5_PROXY", "true", async () => {
    const request = new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        proxy: {
          type: "socks5",
          host: "127.0.0.1",
          port: "1",
        },
      }),
    });

    const response = await proxyTestRoute.POST(request);
    const payload = await response.json();

    assert.notEqual(response.status, 400);
    assert.equal(payload.success, false);
    assert.equal(payload.proxyUrl, "socks5://127.0.0.1:1");
  });
});

test("proxy test route validates JSON, schema, and proxy types before dispatching", async () => {
  await resetStorage();

  const invalidJsonResponse = await proxyTestRoute.POST(
    new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: "{",
    })
  );
  const invalidJsonBody = await invalidJsonResponse.json();
  assert.equal(invalidJsonResponse.status, 400);
  assert.equal(invalidJsonBody.error.message, "Invalid JSON body");

  const invalidBodyResponse = await proxyTestRoute.POST(
    new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({ proxy: { port: "8080" } }),
    })
  );
  const invalidBody = await invalidBodyResponse.json();
  assert.equal(invalidBodyResponse.status, 400);
  assert.equal(invalidBody.error.message, "Invalid request");

  const socks4Response = await proxyTestRoute.POST(
    new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        proxy: {
          type: "socks4",
          host: "127.0.0.1",
          port: "1080",
        },
      }),
    })
  );
  const socks4Body = await socks4Response.json();
  assert.equal(socks4Response.status, 400);
  assert.match(socks4Body.error.message, /proxy\.type must be http or https/i);

  const unsupportedResponse = await proxyTestRoute.POST(
    new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        proxy: {
          type: "ftp",
          host: "127.0.0.1",
          port: "21",
        },
      }),
    })
  );
  const unsupportedBody = await unsupportedResponse.json();
  assert.equal(unsupportedResponse.status, 400);
  assert.match(unsupportedBody.error.message, /proxy\.type must be http or https/i);
});

test("proxy test route handles invalid proxy ports and uses stored proxy config when proxyId is provided", async () => {
  await resetStorage();

  const invalidPortResponse = await proxyTestRoute.POST(
    new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        proxy: {
          type: "http",
          host: "127.0.0.1",
          port: "70000",
        },
      }),
    })
  );
  const invalidPortBody = await invalidPortResponse.json();
  assert.equal(invalidPortResponse.status, 400);
  assert.match(invalidPortBody.error.message, /invalid proxy port/i);

  const storedProxy = await localDb.createProxy({
    name: "Stored Proxy",
    type: "http",
    host: "127.0.0.1",
    port: 1,
    username: "alice",
    password: "secret",
  });

  const proxyIdResponse = await proxyTestRoute.POST(
    new Request("http://localhost/api/settings/proxy/test", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        proxyId: storedProxy.id,
        proxy: {
          host: "127.0.0.1",
          port: 0,
        },
      }),
    })
  );
  const proxyIdBody = await proxyIdResponse.json();
  assert.notEqual(proxyIdResponse.status, 400);
  assert.equal(proxyIdBody.success, false);
  assert.equal(proxyIdBody.proxyUrl, "http://127.0.0.1:1");
});