AgentCompany/tests/unit/provider-rate-limit-registry.test.ts

import { randomUUID } from "node:crypto";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { describe, expect, test } from "vitest";
import { DatabaseConnection } from "../../src/main/database/connection";
import { OrganizationService } from "../../src/main/services/organization-service";
import { parseProviderRateLimitHeaders, ProviderRateLimitRegistry } from "../../src/main/services/provider-rate-limit-registry";
import { SnapshotService } from "../../src/main/services/snapshot-service";

function createConnection(): DatabaseConnection {
  return new DatabaseConnection(join(tmpdir(), `agentcompany-provider-rate-limits-${randomUUID()}.db`));
}

function createOrganization(organization: OrganizationService): string {
  return organization.createOrganization({
    organizationName: "Acme Operations",
    employeeName: "Dana Lee",
    mission: "Coordinate agents with auditable approvals.",
    goal: "Ship a verified control plane.",
  });
}

function providerId(connection: DatabaseConnection, orgId: string, key: string): string {
  return (connection.db.prepare("select id from providers where org_id = ? and provider_key = ?")
    .get(orgId, key) as { id: string }).id;
}

describe("provider rate limit registry", () => {
  test("parses provider quota headers and retry-after hints", () => {
    const now = new Date("2026-05-23T12:00:00.000Z");
    const state = parseProviderRateLimitHeaders(new Headers({
      "x-ratelimit-limit-requests": "100",
      "x-ratelimit-remaining-requests": "42",
      "x-ratelimit-reset-requests": "2s",
      "retry-after": "1.5",
    }), now);

    expect(state).toEqual({
      requestLimit: 100,
      remainingRequests: 42,
      resetAt: "2026-05-23T12:00:02.000Z",
      retryAfterMs: 1500,
      rawStateJson: JSON.stringify({
        "x-ratelimit-limit-requests": "100",
        "x-ratelimit-remaining-requests": "42",
        "x-ratelimit-reset-requests": "2s",
        "retry-after": "1.5",
      }),
    });
  });

  test("persists observed quota state into snapshots", () => {
    const connection = createConnection();
    const organization = new OrganizationService(connection.db);
    const snapshot = new SnapshotService(connection.db);
    const registry = new ProviderRateLimitRegistry(connection.db);
    let orgId = "";
    let openAiProviderId = "";

    connection.transaction(() => {
      orgId = createOrganization(organization);
      openAiProviderId = providerId(connection, orgId, "openai");
      registry.recordObservation({
        orgId,
        providerId: openAiProviderId,
        providerKey: "openai",
        model: "gpt-5.4",
        state: {
          requestLimit: 100,
          remainingRequests: 99,
          resetAt: "2026-05-23T12:10:00.000Z",
          retryAfterMs: null,
          rawStateJson: JSON.stringify({ "x-ratelimit-remaining-requests": "99" }),
        },
      });
    });

    expect(snapshot.load().providerRateLimitObservations[0]).toMatchObject({
      orgId,
      providerId: openAiProviderId,
      providerKey: "openai",
      model: "gpt-5.4",
      requestLimit: 100,
      remainingRequests: 99,
      resetAt: "2026-05-23T12:10:00.000Z",
      retryAfterMs: null,
      rawStateJson: JSON.stringify({ "x-ratelimit-remaining-requests": "99" }),
    });
    connection.close();
  });
});