Add provider rate limit registry
This commit is contained in:
1 parent
9e44cbb31b
commit
d6e797e2b9
16 files changed
+525
-38
No files matched your search
@@ -109,7 +109,12 @@ describe("provider contracts", () => {
|
||||
.get(runId) as { state: string; failed_count: number; result_json: string };
|
||||
const result = JSON.parse(run.result_json) as { checks: Array<{ key: string; passed: boolean }> };
|
||||
expect(run).toMatchObject({ state: SDLC_RUN_STATE.PASSED, failed_count: 0 });
|
||||
expect(result.checks.map((check) => check.key).sort()).toEqual(["provider_adapter_coverage", "provider_runtime_contract", "schema_version"]);
|
||||
expect(result.checks.map((check) => check.key).sort()).toEqual([
|
||||
"provider_adapter_coverage",
|
||||
"provider_rate_limit_registry",
|
||||
"provider_runtime_contract",
|
||||
"schema_version",
|
||||
]);
|
||||
expect(result.checks.every((check) => check.passed)).toBe(true);
|
||||
connection.close();
|
||||
});
|
||||
|
||||
@@ -23,7 +23,7 @@ describe("domain services", () => {
|
||||
const connection = createConnection();
|
||||
|
||||
const rows = connection.db.prepare("select version from schema_migrations order by version asc").all() as Array<{ version: number }>;
|
||||
expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]);
|
||||
expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]);
|
||||
connection.close();
|
||||
});
|
||||
|
||||
|
||||
@@ -132,6 +132,7 @@ describe("migration service", () => {
|
||||
create table providers(id text primary key, org_id text, credential_id text);
|
||||
create table provider_requests(id text primary key, org_id text, provider_id text);
|
||||
create table provider_usage_windows(id text primary key, org_id text, provider_id text);
|
||||
create table provider_rate_limit_observations(id text primary key, org_id text, provider_id text);
|
||||
create table credentials(id text primary key, org_id text);
|
||||
create table skill_definitions(id text primary key, org_id text);
|
||||
create table skill_revisions(id text primary key, skill_id text);
|
||||
@@ -199,6 +200,7 @@ describe("migration service", () => {
|
||||
insert into chat_read_state(org_id, employee_id, space_id, last_read_message_id) values ('org-1', 'emp-1', 'space-1', 'msg-1');
|
||||
insert into provider_requests(id, org_id, provider_id) values ('req-1', 'org-1', 'prov-1');
|
||||
insert into provider_usage_windows(id, org_id, provider_id) values ('usage-1', 'org-1', 'prov-1');
|
||||
insert into provider_rate_limit_observations(id, org_id, provider_id) values ('quota-1', 'org-1', 'prov-1');
|
||||
insert into skill_definitions(id, org_id) values ('skill-1', 'org-1');
|
||||
insert into skill_revisions(id, skill_id) values ('rev-1', 'skill-1'), ('rev-2', 'skill-1');
|
||||
insert into skill_artifacts(id, revision_id) values ('artifact-1', 'rev-1');
|
||||
@@ -244,6 +246,7 @@ describe("migration service", () => {
|
||||
expect(body.tableCounts.file_access_grants).toBe(1);
|
||||
expect(body.tableCounts.provider_requests).toBe(1);
|
||||
expect(body.tableCounts.provider_usage_windows).toBe(1);
|
||||
expect(body.tableCounts.provider_rate_limit_observations).toBe(1);
|
||||
expect(body.tableCounts.compliance_exports).toBe(1);
|
||||
expect(body.tableCounts.audit_replay_runs).toBe(1);
|
||||
expect(body.tableCounts.rollback_snapshots).toBe(1);
|
||||
@@ -262,6 +265,7 @@ describe("migration service", () => {
|
||||
expect(body.idMaps.audit_replay_runs).toEqual(["replay-1"]);
|
||||
expect(body.idMaps.rollback_snapshots).toEqual(["rollback-1"]);
|
||||
expect(body.idMaps.provider_usage_windows).toEqual(["usage-1"]);
|
||||
expect(body.idMaps.provider_rate_limit_observations).toEqual(["quota-1"]);
|
||||
connection.close();
|
||||
});
|
||||
|
||||
|
||||
@@ -64,8 +64,8 @@ describe("operations service", () => {
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString() }),
|
||||
createResponse: async () => ({ text: "accepted", rawResponseJson: "{}" }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
|
||||
createResponse: async () => ({ text: "accepted", rawResponseJson: "{}", rateLimitState: null }),
|
||||
},
|
||||
};
|
||||
const providers = new ProviderService(connection.db, adapters, testVault);
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
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();
|
||||
});
|
||||
});
|
||||
@@ -5,7 +5,7 @@ import { describe, expect, test } from "vitest";
|
||||
import { DatabaseConnection } from "../../src/main/database/connection";
|
||||
import { OrganizationService } from "../../src/main/services/organization-service";
|
||||
import { ProviderService } from "../../src/main/services/provider-service";
|
||||
import { providerAdapters, type ProviderAdapterRegistry, type ProviderAuthProfile } from "../../src/main/services/provider-adapters";
|
||||
import { ProviderHttpError, providerAdapters, type ProviderAdapterRegistry, type ProviderAuthProfile } from "../../src/main/services/provider-adapters";
|
||||
import { SnapshotService } from "../../src/main/services/snapshot-service";
|
||||
import { PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE } from "../../src/shared/domain";
|
||||
import { PROVIDER_CATALOG } from "../../src/shared/provider-catalog";
|
||||
@@ -144,7 +144,7 @@ describe("provider service", () => {
|
||||
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
|
||||
checkHealth: async ({ authProfile }) => {
|
||||
checkedProfile = authProfile;
|
||||
return { checkedAt: new Date().toISOString() };
|
||||
return { checkedAt: new Date().toISOString(), rateLimitState: null };
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -178,10 +178,11 @@ describe("provider service", () => {
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
|
||||
createResponse: async ({ authProfile, model, prompt }) => ({
|
||||
text: `${authProfile.type}:${model}:${prompt}`,
|
||||
rawResponseJson: JSON.stringify({ provider: "openai" }),
|
||||
rateLimitState: null,
|
||||
}),
|
||||
},
|
||||
};
|
||||
@@ -214,6 +215,116 @@ describe("provider service", () => {
|
||||
connection.close();
|
||||
});
|
||||
|
||||
test("records provider response rate limit observations", async () => {
|
||||
const connection = createConnection();
|
||||
const organization = new OrganizationService(connection.db);
|
||||
const snapshot = new SnapshotService(connection.db);
|
||||
const adapters: ProviderAdapterRegistry = {
|
||||
...providerAdapters,
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
|
||||
createResponse: async () => ({
|
||||
text: "accepted",
|
||||
rawResponseJson: "{}",
|
||||
rateLimitState: {
|
||||
requestLimit: 120,
|
||||
remainingRequests: 119,
|
||||
resetAt: "2026-05-23T12:00:00.000Z",
|
||||
retryAfterMs: null,
|
||||
rawStateJson: JSON.stringify({ "x-ratelimit-remaining-requests": "119" }),
|
||||
},
|
||||
}),
|
||||
},
|
||||
};
|
||||
const providers = new ProviderService(connection.db, adapters, testVault);
|
||||
let orgId = "";
|
||||
let openAiProviderId = "";
|
||||
|
||||
connection.transaction(() => {
|
||||
orgId = createOrganization(organization);
|
||||
openAiProviderId = providerId(connection, orgId, "openai");
|
||||
providers.saveCredential({ orgId, providerId: openAiProviderId, secret: "sk-live-provider", model: "gpt-5.4" });
|
||||
});
|
||||
await providers.createResponse({
|
||||
orgId,
|
||||
providerId: openAiProviderId,
|
||||
prompt: "Draft board memo",
|
||||
instructions: null,
|
||||
});
|
||||
|
||||
expect(snapshot.load().providerRateLimitObservations[0]).toMatchObject({
|
||||
orgId,
|
||||
providerId: openAiProviderId,
|
||||
providerKey: "openai",
|
||||
model: "gpt-5.4",
|
||||
requestLimit: 120,
|
||||
remainingRequests: 119,
|
||||
resetAt: "2026-05-23T12:00:00.000Z",
|
||||
retryAfterMs: null,
|
||||
rawStateJson: JSON.stringify({ "x-ratelimit-remaining-requests": "119" }),
|
||||
});
|
||||
connection.close();
|
||||
});
|
||||
|
||||
test("records failed provider response rate limit observations", async () => {
|
||||
const connection = createConnection();
|
||||
const organization = new OrganizationService(connection.db);
|
||||
const snapshot = new SnapshotService(connection.db);
|
||||
const adapters: ProviderAdapterRegistry = {
|
||||
...providerAdapters,
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
|
||||
createResponse: async () => {
|
||||
throw new ProviderHttpError("rate limited", {
|
||||
requestLimit: 120,
|
||||
remainingRequests: 0,
|
||||
resetAt: "2026-05-23T12:00:00.000Z",
|
||||
retryAfterMs: 2000,
|
||||
rawStateJson: JSON.stringify({ "retry-after": "2" }),
|
||||
});
|
||||
},
|
||||
},
|
||||
};
|
||||
const providers = new ProviderService(connection.db, adapters, testVault);
|
||||
let orgId = "";
|
||||
let openAiProviderId = "";
|
||||
|
||||
connection.transaction(() => {
|
||||
orgId = createOrganization(organization);
|
||||
openAiProviderId = providerId(connection, orgId, "openai");
|
||||
providers.saveCredential({ orgId, providerId: openAiProviderId, secret: "sk-live-provider", model: "gpt-5.4" });
|
||||
});
|
||||
await expect(providers.createResponse({
|
||||
orgId,
|
||||
providerId: openAiProviderId,
|
||||
prompt: "Draft board memo",
|
||||
instructions: null,
|
||||
})).rejects.toThrow("rate limited");
|
||||
|
||||
const state = snapshot.load();
|
||||
expect(state.providerRequests[0]).toMatchObject({
|
||||
providerId: openAiProviderId,
|
||||
state: "failed",
|
||||
errorMessage: "rate limited",
|
||||
});
|
||||
expect(state.providerRateLimitObservations[0]).toMatchObject({
|
||||
orgId,
|
||||
providerId: openAiProviderId,
|
||||
providerKey: "openai",
|
||||
model: "gpt-5.4",
|
||||
requestLimit: 120,
|
||||
remainingRequests: 0,
|
||||
resetAt: "2026-05-23T12:00:00.000Z",
|
||||
retryAfterMs: 2000,
|
||||
rawStateJson: JSON.stringify({ "retry-after": "2" }),
|
||||
});
|
||||
connection.close();
|
||||
});
|
||||
|
||||
test("tracks provider usage windows and enforces monthly request limits", async () => {
|
||||
const connection = createConnection();
|
||||
const organization = new OrganizationService(connection.db);
|
||||
@@ -223,8 +334,8 @@ describe("provider service", () => {
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString() }),
|
||||
createResponse: async ({ prompt }) => ({ text: `accepted:${prompt}`, rawResponseJson: "{}" }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
|
||||
createResponse: async ({ prompt }) => ({ text: `accepted:${prompt}`, rawResponseJson: "{}", rateLimitState: null }),
|
||||
},
|
||||
};
|
||||
const providers = new ProviderService(connection.db, adapters, testVault);
|
||||
@@ -292,10 +403,10 @@ describe("provider service", () => {
|
||||
openai: {
|
||||
...providerAdapters.openai,
|
||||
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString() }),
|
||||
checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
|
||||
createResponse: async () => {
|
||||
adapterCalled = true;
|
||||
return { text: "blocked", rawResponseJson: "{}" };
|
||||
return { text: "blocked", rawResponseJson: "{}", rateLimitState: null };
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
Reference in new issue
Block a user