OmniRoute/tests/unit/combo-selected-connection-success.test.ts

/* eslint-disable @typescript-eslint/no-explicit-any */
import test, { describe, beforeEach } from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";

const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-combo-sel-conn-"));
const ORIGINAL_DATA_DIR = process.env.DATA_DIR;
process.env.DATA_DIR = TEST_DATA_DIR;

const { handleComboChat } = await import("../../open-sse/services/combo.ts");
const {
  recordModelLockoutFailure,
  getModelLockoutInfo,
  clearAllModelLockouts,
  decayModelFailureCount,
} = await import("../../open-sse/services/accountFallback.ts");
const { recordProviderCooldown, isProviderInCooldown, recordProviderSuccess, clearCooldownState } =
  await import("../../open-sse/services/providerCooldownTracker.ts");
const { resolveResilienceSettings } = await import("../../src/lib/resilience/settings.ts");
const core = await import("../../src/lib/db/core.ts");
const settingsDb = await import("../../src/lib/db/settings.ts");
const providersDb = await import("../../src/lib/db/providers.ts");

const settings = resolveResilienceSettings({
  resilienceSettings: {
    providerCooldown: {
      enabled: true,
      minRetryCooldownMs: 5000,
      maxRetryCooldownMs: 300000,
    },
  },
});

function createLog() {
  return {
    info: (tag: any, msg: any) => console.log(`[INFO][${tag}] ${msg}`),
    warn: (tag: any, msg: any) => console.log(`[WARN][${tag}] ${msg}`),
    error: (tag: any, msg: any) => console.log(`[ERROR][${tag}] ${msg}`),
    debug: (tag: any, msg: any) => console.log(`[DEBUG][${tag}] ${msg}`),
  };
}

async function cleanupTestDataDir() {
  let lastError;
  for (let attempt = 0; attempt < 5; attempt += 1) {
    try {
      core.resetDbInstance();
      fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 });
      return;
    } catch (error: any) {
      lastError = error;
      await new Promise((resolve) => setTimeout(resolve, 25));
    }
  }
  if (lastError) {
    throw lastError;
  }
}

test.after(async () => {
  await cleanupTestDataDir();
  process.env.DATA_DIR = ORIGINAL_DATA_DIR;
});

beforeEach(async () => {
  clearAllModelLockouts();
  clearCooldownState();
  settingsDb.clearAllLKGP();
});

describe("combo selected connection success handling", () => {
  test("priority strategy correctly extracts dynamic connection ID from success response headers and decays lockout, resets provider cooldown, and updates LKGP", async () => {
    const comboName = "test-combo-priority";
    const modelStr = "openai/gpt-4";
    const provider = "openai";
    const dynamicConnId = "conn-dynamic-123";

    await providersDb.createProviderConnection({
      provider,
      authType: "apikey",
      name: "OpenAI Test",
      apiKey: "sk-test",
    });

    // 1. Populate lockout failure count = 4
    recordModelLockoutFailure(
      provider,
      dynamicConnId,
      "gpt-4",
      "rate_limit_exceeded",
      429,
      120_000,
      null,
      { exactCooldownMs: 60_000 }
    );
    for (let i = 0; i < 3; i++) {
      recordModelLockoutFailure(
        provider,
        dynamicConnId,
        "gpt-4",
        "rate_limit_exceeded",
        429,
        120_000,
        null,
        { exactCooldownMs: 60_000 }
      );
    }

    // Verify initial failureCount is 4
    const initialLockout = getModelLockoutInfo(provider, dynamicConnId, "gpt-4");
    assert.equal(initialLockout?.failureCount, 4);

    // 2. Record provider cooldown
    recordProviderCooldown(provider, dynamicConnId, settings);
    assert.ok(isProviderInCooldown(provider, dynamicConnId, settings));

    // 3. Invoke handleComboChat
    const result = await handleComboChat({
      body: { stream: false },
      combo: {
        name: comboName,
        strategy: "priority",
        models: [modelStr],
        config: { maxRetries: 0, concurrencyPerModel: 1, queueTimeoutMs: 1000 },
      },
      handleSingleModel: async () => {
        return new Response(JSON.stringify({ choices: [{ message: { content: "ok" } }] }), {
          status: 200,
          headers: {
            "content-type": "application/json",
            "X-OmniRoute-Selected-Connection-Id": dynamicConnId,
          },
        });
      },
      isModelAvailable: async () => true,
      log: createLog(),
      settings: null,
      allCombos: null,
    });

    assert.equal(result.ok, true);

    // 4. Assertions
    // A. Dynamic connection-level failure count decay (halved to 2)
    const decayCheck = decayModelFailureCount(provider, dynamicConnId, "gpt-4");
    assert.equal(
      decayCheck.newFailureCount,
      1,
      "failure count should decay from 4 to 2, and now to 1"
    );

    // B. Dynamic connection-level provider success tracking (not in cooldown anymore)
    assert.equal(
      isProviderInCooldown(provider, dynamicConnId, settings),
      false,
      "provider cooldown should be cleared on success"
    );

    // C. Correct LKGP record updated with the dynamic connection ID (setLKGP is called)
    let persisted: any = null;
    for (let i = 0; i < 20; i++) {
      persisted = await settingsDb.getLKGP(comboName, comboName);
      if (persisted?.connectionId === dynamicConnId) {
        break;
      }
      await new Promise((resolve) => setTimeout(resolve, 10));
    }
    assert.equal(persisted?.provider, provider);
    assert.equal(
      persisted?.connectionId,
      dynamicConnId,
      "LKGP connectionId must be the dynamic connection ID"
    );
  });

  test("priority strategy with lowercase selected connection ID header correctly decays lockout, resets provider cooldown, and updates LKGP", async () => {
    const comboName = "test-combo-priority-lc";
    const modelStr = "openai/gpt-4";
    const provider = "openai";
    const dynamicConnId = "conn-dynamic-123-lc";

    await providersDb.createProviderConnection({
      provider,
      authType: "apikey",
      name: "OpenAI Test LC",
      apiKey: "sk-test",
    });

    // 1. Populate lockout failure count = 4
    recordModelLockoutFailure(
      provider,
      dynamicConnId,
      "gpt-4",
      "rate_limit_exceeded",
      429,
      120_000,
      null,
      { exactCooldownMs: 60_000 }
    );
    for (let i = 0; i < 3; i++) {
      recordModelLockoutFailure(
        provider,
        dynamicConnId,
        "gpt-4",
        "rate_limit_exceeded",
        429,
        120_000,
        null,
        { exactCooldownMs: 60_000 }
      );
    }

    // Verify initial failureCount is 4
    const initialLockout = getModelLockoutInfo(provider, dynamicConnId, "gpt-4");
    assert.equal(initialLockout?.failureCount, 4);

    // 2. Record provider cooldown
    recordProviderCooldown(provider, dynamicConnId, settings);
    assert.ok(isProviderInCooldown(provider, dynamicConnId, settings));

    // 3. Invoke handleComboChat
    const result = await handleComboChat({
      body: { stream: false },
      combo: {
        name: comboName,
        strategy: "priority",
        models: [modelStr],
        config: { maxRetries: 0, concurrencyPerModel: 1, queueTimeoutMs: 1000 },
      },
      handleSingleModel: async () => {
        return new Response(JSON.stringify({ choices: [{ message: { content: "ok" } }] }), {
          status: 200,
          headers: {
            "content-type": "application/json",
            "x-omniroute-selected-connection-id": dynamicConnId,
          },
        });
      },
      isModelAvailable: async () => true,
      log: createLog(),
      settings: null,
      allCombos: null,
    });

    assert.equal(result.ok, true);

    // 4. Assertions
    const decayCheck = decayModelFailureCount(provider, dynamicConnId, "gpt-4");
    assert.equal(
      decayCheck.newFailureCount,
      1,
      "failure count should decay from 4 to 2, and now to 1"
    );

    assert.equal(
      isProviderInCooldown(provider, dynamicConnId, settings),
      false,
      "provider cooldown should be cleared on success"
    );

    let persisted: any = null;
    for (let i = 0; i < 20; i++) {
      persisted = await settingsDb.getLKGP(comboName, comboName);
      if (persisted?.connectionId === dynamicConnId) {
        break;
      }
      await new Promise((resolve) => setTimeout(resolve, 10));
    }
    assert.equal(persisted?.provider, provider);
    assert.equal(
      persisted?.connectionId,
      dynamicConnId,
      "LKGP connectionId must be the dynamic connection ID"
    );
  });

  test("round-robin strategy correctly extracts dynamic connection ID from success response headers and decays lockout, resets provider cooldown, and updates LKGP", async () => {
    const comboName = "test-combo-rr";
    const modelStr = "openai/gpt-4";
    const provider = "openai";
    const dynamicConnId = "conn-dynamic-123-rr";

    await providersDb.createProviderConnection({
      provider,
      authType: "apikey",
      name: "OpenAI Test RR",
      apiKey: "sk-test",
    });

    // 1. Populate lockout failure count = 4
    recordModelLockoutFailure(
      provider,
      dynamicConnId,
      "gpt-4",
      "rate_limit_exceeded",
      429,
      120_000,
      null,
      { exactCooldownMs: 60_000 }
    );
    for (let i = 0; i < 3; i++) {
      recordModelLockoutFailure(
        provider,
        dynamicConnId,
        "gpt-4",
        "rate_limit_exceeded",
        429,
        120_000,
        null,
        { exactCooldownMs: 60_000 }
      );
    }

    // Verify initial failureCount is 4
    const initialLockout = getModelLockoutInfo(provider, dynamicConnId, "gpt-4");
    assert.equal(initialLockout?.failureCount, 4);

    // 2. Record provider cooldown
    recordProviderCooldown(provider, dynamicConnId, settings);
    assert.ok(isProviderInCooldown(provider, dynamicConnId, settings));

    // 3. Invoke handleComboChat with round-robin strategy
    const result = await handleComboChat({
      body: { stream: false },
      combo: {
        name: comboName,
        strategy: "round-robin",
        models: [modelStr],
        config: { maxRetries: 0, concurrencyPerModel: 1, queueTimeoutMs: 1000 },
      },
      handleSingleModel: async () => {
        return new Response(JSON.stringify({ choices: [{ message: { content: "ok" } }] }), {
          status: 200,
          headers: {
            "content-type": "application/json",
            "X-OmniRoute-Selected-Connection-Id": dynamicConnId,
          },
        });
      },
      isModelAvailable: async () => true,
      log: createLog(),
      settings: null,
      allCombos: null,
    });

    assert.equal(result.ok, true);

    // 4. Assertions
    const decayCheck = decayModelFailureCount(provider, dynamicConnId, "gpt-4");
    assert.equal(
      decayCheck.newFailureCount,
      1,
      "failure count should decay from 4 to 2, and now to 1"
    );

    assert.equal(
      isProviderInCooldown(provider, dynamicConnId, settings),
      false,
      "provider cooldown should be cleared on success"
    );

    let persisted: any = null;
    for (let i = 0; i < 20; i++) {
      persisted = await settingsDb.getLKGP(comboName, comboName);
      if (persisted?.connectionId === dynamicConnId) {
        break;
      }
      await new Promise((resolve) => setTimeout(resolve, 10));
    }
    assert.equal(persisted?.provider, provider);
    assert.equal(
      persisted?.connectionId,
      dynamicConnId,
      "LKGP connectionId must be the dynamic connection ID"
    );
  });
});