qinglong/packages/ql3-ai/test/postgresModelProviderCredentialTestConnection.test.cjs

const assert = require('node:assert/strict');
const { randomUUID } = require('node:crypto');
const test = require('node:test');

const {
  createModelProviderCredentialTestAllowlist,
  createModelProviderCredentialTestPlan,
} = require('../dist/model-provider-credential/modelProviderCredentialTestConnection.js');
const {
  MODEL_PROVIDER_CREDENTIAL_TEST_PLAN_OPERATION_ID,
  ModelProviderCredentialTestPlanAuthorizationFenceConflictError,
  ModelProviderCredentialTestPlanQuotaExceededError,
  ModelProviderCredentialTestPlanUnavailableError,
  PostgresModelProviderCredentialTestPlanRepository,
} = require('../dist/model-provider-credential/postgresModelProviderCredentialTestConnection.js');

function plan(overrides = {}) {
  const endpoint = createModelProviderCredentialTestAllowlist({
    revision: 'catalog-v1',
    providers: [
      {
        provider: 'openai-compatible',
        adapter: 'openai-compatible',
        baseUrl: 'https://provider.example.test/v1/',
        revision: 'endpoint-v1',
        deadlineMs: 5_000,
        maxResponseBytes: 64 * 1_024,
        maxModels: 64,
        maxCostMicrousd: 0,
        retryLimit: 0,
      },
    ],
  }).providers[0];
  return createModelProviderCredentialTestPlan({
    testId: randomUUID(),
    requestId: `request-${randomUUID()}`,
    projectId: 'project-a',
    provider: endpoint.provider,
    endpoint,
    requestedBy: { type: 'user', id: 'owner-a' },
    fence: { projectVersion: 1, bindingVersion: 1 },
    plannedAtMs: 100,
    expiresAtMs: 60_100,
    ...overrides,
  });
}

function authorized(value) {
  return {
    plan: value,
    audit: {
      eventId: value.testId,
      requestId: value.requestId,
      operationId: MODEL_PROVIDER_CREDENTIAL_TEST_PLAN_OPERATION_ID,
      projectId: value.projectId,
      subject: value.requestedBy,
      authenticationId: 'strong-authentication-1',
      outcome: 'allowed',
      reasons: ['project_owner'],
      fence: value.fence,
      occurredAtMs: value.plannedAtMs,
    },
  };
}

function fixture({ projectVersion = 1, bindingVersion = 1 } = {}) {
  const state = {
    plans: new Map(),
    audits: new Map(),
    quota: { consumed: 0, receipts: new Set() },
  };
  let snapshot;
  let loseCommit = false;
  const client = {
    async query(text, values = []) {
      if (text.startsWith('BEGIN')) {
        snapshot = structuredClone(state);
        return { rows: [] };
      }
      if (text.includes('pg_advisory_xact_lock')) return { rows: [{}] };
      if (text.includes('FROM "ql3"."projects"')) {
        return { rows: [{ status: 'active', version: projectVersion }] };
      }
      if (text.includes('FROM "ql3"."project_role_bindings"')) {
        return { rows: [{ state: 'active', version: bindingVersion }] };
      }
      if (
        text.includes('FROM "ql3_ai"."model_provider_credential_test_plans"')
      ) {
        const [testId, requestedProjectId, requestId] = values;
        const rows = [...state.plans.values()]
          .filter(
            (stored) =>
              stored.testId === testId ||
              (stored.projectId === requestedProjectId &&
                stored.requestId === requestId),
          )
          .map((stored) => ({ planJson: stored }));
        return { rows };
      }
      if (text.includes('FROM "ql3"."security_audit_events"')) {
        const stored = state.audits.get(values[0]);
        return { rows: stored ? [stored] : [] };
      }
      if (text.includes('receipt_ids ? $3::text AS "hasReceipt"')) {
        return {
          rows:
            state.quota.consumed === 0
              ? []
              : [{ hasReceipt: state.quota.receipts.has(values[2]) }],
        };
      }
      if (
        text.startsWith('WITH database_clock AS (') &&
        text.includes('INSERT INTO')
      ) {
        const receipt = values[2];
        const limit = values[4];
        if (
          !state.quota.receipts.has(receipt) &&
          state.quota.consumed >= limit
        ) {
          return { rows: [] };
        }
        if (!state.quota.receipts.has(receipt)) {
          state.quota.receipts.add(receipt);
          state.quota.consumed += 1;
        }
        return {
          rows: [
            {
              consumedCount: state.quota.consumed,
              resetAtMs: 60_100,
              observedAtMs: 100,
            },
          ],
        };
      }
      if (
        text.startsWith('WITH database_clock AS (') &&
        text.includes('SELECT consumed_count')
      ) {
        return {
          rows: [
            {
              consumedCount: state.quota.consumed,
              resetAtMs: 60_100,
              observedAtMs: 100,
            },
          ],
        };
      }
      if (
        text.startsWith(
          'INSERT INTO "ql3_ai"."model_provider_credential_test_plans"',
        )
      ) {
        state.plans.set(values[0], JSON.parse(values[20]));
        return { rows: [] };
      }
      if (text.startsWith('INSERT INTO "ql3"."security_audit_events"')) {
        state.audits.set(values[0], {
          eventId: values[0],
          requestId: values[1],
          operationId: values[2],
          projectId: values[3],
          subjectType: values[4],
          subjectId: values[5],
          authenticationId: values[6],
          outcome: values[7],
          reasons: JSON.parse(values[8]),
          projectVersion: values[9],
          bindingVersion: values[10],
          occurredAtMs: values[11],
        });
        return { rows: [] };
      }
      if (text === 'COMMIT') {
        snapshot = undefined;
        if (loseCommit) {
          loseCommit = false;
          const error = new Error('lost commit response');
          error.code = 'ECONNRESET';
          throw error;
        }
        return { rows: [] };
      }
      if (text === 'ROLLBACK') {
        if (snapshot) {
          state.plans = snapshot.plans;
          state.audits = snapshot.audits;
          state.quota = snapshot.quota;
        }
        snapshot = undefined;
        return { rows: [] };
      }
      throw new Error(`unexpected query: ${text}`);
    },
    release() {},
  };
  return {
    repository: new PostgresModelProviderCredentialTestPlanRepository(
      {
        async connect() {
          return client;
        },
      },
      { quotaWindowMs: 60_000, quotaLimit: 2 },
    ),
    state: () => state,
    loseNextCommitResponse() {
      loseCommit = true;
    },
  };
}

test('atomically consumes quota, stores a plan and writes allowed audit', async () => {
  const value = fixture();
  const candidate = plan();
  const created = await value.repository.createAuthorized(
    authorized(candidate),
  );
  assert.equal(created.status, 'created');
  assert.equal(value.state().plans.size, 1);
  assert.equal(value.state().audits.size, 1);
  assert.equal(value.state().quota.consumed, 1);
  assert.equal(value.state().quota.receipts.has(candidate.testId), true);
});

test('converges a COMMIT response loss without consuming quota twice', async () => {
  const value = fixture();
  const candidate = plan();
  value.loseNextCommitResponse();
  await assert.rejects(
    value.repository.createAuthorized(authorized(candidate)),
    ModelProviderCredentialTestPlanUnavailableError,
  );
  const replay = await value.repository.createAuthorized(authorized(candidate));
  assert.equal(replay.status, 'existing');
  assert.equal(value.state().quota.consumed, 1);
});

test('replays the stored plan when API retry observes a later clock', async () => {
  const value = fixture();
  const candidate = plan();
  await value.repository.createAuthorized(authorized(candidate));
  const retried = plan({
    testId: candidate.testId,
    requestId: candidate.requestId,
    projectId: candidate.projectId,
    provider: candidate.provider,
    endpoint: candidate.endpoint,
    requestedBy: candidate.requestedBy,
    fence: candidate.fence,
    plannedAtMs: candidate.plannedAtMs + 25,
    expiresAtMs: candidate.expiresAtMs + 25,
  });
  const replay = await value.repository.createAuthorized(authorized(retried));
  assert.equal(replay.status, 'existing');
  assert.deepEqual(replay.plan, candidate);
  assert.equal(value.state().quota.consumed, 1);
});

test('fails a stale fence and quota excess before creating another plan', async () => {
  const stale = fixture({ projectVersion: 2 });
  await assert.rejects(
    stale.repository.createAuthorized(authorized(plan())),
    ModelProviderCredentialTestPlanAuthorizationFenceConflictError,
  );
  assert.equal(stale.state().quota.consumed, 0);

  const limited = fixture();
  await limited.repository.createAuthorized(authorized(plan()));
  await limited.repository.createAuthorized(authorized(plan()));
  await assert.rejects(
    limited.repository.createAuthorized(authorized(plan())),
    ModelProviderCredentialTestPlanQuotaExceededError,
  );
  assert.equal(limited.state().plans.size, 2);
});