Add provider rate limit registry

This commit is contained in:
ZacharyZhang-NY committed 2026-05-23 19:50:20 -04:00
1 parent 9e44cbb31b
commit d6e797e2b9
16 files changed
+525 -38

No files matched your search

@@ -0,0 +1,21 @@
export const providerRateLimitObservationSchema = `
create table if not exists provider_rate_limit_observations (
id text primary key,
org_id text not null references organizations(id) on delete cascade,
provider_id text not null references providers(id) on delete cascade,
provider_key text not null,
model text,
request_limit integer,
remaining_requests integer,
reset_at text,
retry_after_ms integer,
raw_state_json text not null,
observed_at text not null
);
create index if not exists idx_provider_rate_limit_observations_org_provider
on provider_rate_limit_observations(org_id, provider_id, observed_at desc);
insert into schema_migrations(version, name, applied_at)
values (22, 'provider_rate_limit_observations', datetime('now'));
`;
@@ -20,6 +20,7 @@ import { sdlcControlsSchema } from "./migrations/018-sdlc-controls";
import { auditReplaySchema } from "./migrations/019-audit-replay"; import { auditReplaySchema } from "./migrations/019-audit-replay";
import { rollbackSnapshotsSchema } from "./migrations/020-rollback-snapshots"; import { rollbackSnapshotsSchema } from "./migrations/020-rollback-snapshots";
import { providerAuthBrokerSchema } from "./migrations/021-provider-auth-broker"; import { providerAuthBrokerSchema } from "./migrations/021-provider-auth-broker";
import { providerRateLimitObservationSchema } from "./migrations/022-provider-rate-limit-observations";
interface Migration { interface Migration {
version: number; version: number;
@@ -49,6 +50,7 @@ const MIGRATIONS: Migration[] = [
{ version: 19, name: "audit_replay", sql: auditReplaySchema }, { version: 19, name: "audit_replay", sql: auditReplaySchema },
{ version: 20, name: "rollback_snapshots", sql: rollbackSnapshotsSchema }, { version: 20, name: "rollback_snapshots", sql: rollbackSnapshotsSchema },
{ version: 21, name: "provider_auth_broker", sql: providerAuthBrokerSchema }, { version: 21, name: "provider_auth_broker", sql: providerAuthBrokerSchema },
{ version: 22, name: "provider_rate_limit_observations", sql: providerRateLimitObservationSchema },
]; ];
const SCHEMA_VERSION = MIGRATIONS[MIGRATIONS.length - 1].version; const SCHEMA_VERSION = MIGRATIONS[MIGRATIONS.length - 1].version;
@@ -50,6 +50,7 @@ const LEGACY_TABLES = [
"providers", "providers",
"provider_requests", "provider_requests",
"provider_usage_windows", "provider_usage_windows",
"provider_rate_limit_observations",
"credentials", "credentials",
"skill_definitions", "skill_definitions",
"skill_revisions", "skill_revisions",
@@ -185,6 +186,8 @@ const RELATION_CHECKS: RelationCheck[] = [
{ table: "provider_requests", column: "provider_id", referencedTable: "providers", referencedColumn: "id", reason: "dangling_provider_id" }, { table: "provider_requests", column: "provider_id", referencedTable: "providers", referencedColumn: "id", reason: "dangling_provider_id" },
{ table: "provider_usage_windows", column: "org_id", referencedTable: "organizations", referencedColumn: "id", reason: "dangling_org_id" }, { table: "provider_usage_windows", column: "org_id", referencedTable: "organizations", referencedColumn: "id", reason: "dangling_org_id" },
{ table: "provider_usage_windows", column: "provider_id", referencedTable: "providers", referencedColumn: "id", reason: "dangling_provider_id" }, { table: "provider_usage_windows", column: "provider_id", referencedTable: "providers", referencedColumn: "id", reason: "dangling_provider_id" },
{ table: "provider_rate_limit_observations", column: "org_id", referencedTable: "organizations", referencedColumn: "id", reason: "dangling_org_id" },
{ table: "provider_rate_limit_observations", column: "provider_id", referencedTable: "providers", referencedColumn: "id", reason: "dangling_provider_id" },
{ table: "credentials", column: "org_id", referencedTable: "organizations", referencedColumn: "id", reason: "dangling_org_id" }, { table: "credentials", column: "org_id", referencedTable: "organizations", referencedColumn: "id", reason: "dangling_org_id" },
{ table: "skill_definitions", column: "org_id", referencedTable: "organizations", referencedColumn: "id", reason: "dangling_org_id" }, { table: "skill_definitions", column: "org_id", referencedTable: "organizations", referencedColumn: "id", reason: "dangling_org_id" },
{ table: "skill_revisions", column: "skill_id", referencedTable: "skill_definitions", referencedColumn: "id", reason: "dangling_skill_id" }, { table: "skill_revisions", column: "skill_id", referencedTable: "skill_definitions", referencedColumn: "id", reason: "dangling_skill_id" },
@@ -1,5 +1,7 @@
import { PROVIDER_AUTH_TYPE, type ProviderAuthType } from "../../shared/domain"; import { PROVIDER_AUTH_TYPE, type ProviderAuthType } from "../../shared/domain";
import { listProviderModels, type ProviderModelEntry } from "../../shared/provider-models"; import { listProviderModels, type ProviderModelEntry } from "../../shared/provider-models";
import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain";
import { parseProviderRateLimitHeaders } from "./provider-rate-limit-registry";
export interface ProviderExecutionInput { export interface ProviderExecutionInput {
authProfile: ProviderAuthProfile; authProfile: ProviderAuthProfile;
@@ -11,6 +13,7 @@ export interface ProviderExecutionInput {
export interface ProviderExecutionResult { export interface ProviderExecutionResult {
text: string; text: string;
rawResponseJson: string; rawResponseJson: string;
rateLimitState: ProviderObservedRateLimitState | null;
} }
export interface ProviderHealthInput { export interface ProviderHealthInput {
@@ -20,6 +23,7 @@ export interface ProviderHealthInput {
export interface ProviderHealthResult { export interface ProviderHealthResult {
checkedAt: string; checkedAt: string;
rateLimitState: ProviderObservedRateLimitState | null;
} }
export interface ProviderContext { export interface ProviderContext {
@@ -31,6 +35,7 @@ export interface ProviderRateLimitState {
requestLimit: number | null; requestLimit: number | null;
remainingRequests: number | null; remainingRequests: number | null;
resetAt: string | null; resetAt: string | null;
retryAfterMs: number | null;
rawStateJson: string; rawStateJson: string;
} }
@@ -58,6 +63,25 @@ type JsonRecord = Record<string, unknown>;
const PROVIDER_REQUEST_TIMEOUT_MS = 60_000; const PROVIDER_REQUEST_TIMEOUT_MS = 60_000;
interface ProviderJsonResponse {
data: JsonRecord;
rateLimitState: ProviderObservedRateLimitState | null;
}
export class ProviderHttpError extends Error {
constructor(
message: string,
readonly rateLimitState: ProviderObservedRateLimitState | null,
) {
super(message);
this.name = "ProviderHttpError";
}
}
export function providerRateLimitStateFromError(error: unknown): ProviderObservedRateLimitState | null {
return error instanceof ProviderHttpError ? error.rateLimitState : null;
}
export const providerAdapters: ProviderAdapterRegistry = { export const providerAdapters: ProviderAdapterRegistry = {
openai: { openai: {
providerKey: "openai", providerKey: "openai",
@@ -68,15 +92,16 @@ export const providerAdapters: ProviderAdapterRegistry = {
checkHealth: (input) => checkOpenAiHealth(input.authProfile, "https://api.openai.com/v1/models"), checkHealth: (input) => checkOpenAiHealth(input.authProfile, "https://api.openai.com/v1/models"),
async createResponse(input) { async createResponse(input) {
const profile = requireApiKeyProfile(input.authProfile, "OpenAI"); const profile = requireApiKeyProfile(input.authProfile, "OpenAI");
const data = await postJson("https://api.openai.com/v1/responses", { const response = await postJson("https://api.openai.com/v1/responses", {
authorization: `Bearer ${profile.apiKey}`, authorization: `Bearer ${profile.apiKey}`,
}, { }, {
model: input.model, model: input.model,
input: input.prompt, input: input.prompt,
...(input.instructions ? { instructions: input.instructions } : {}), ...(input.instructions ? { instructions: input.instructions } : {}),
}); });
const data = response.data;
const text = requireText(collectText(data.output_text ?? data.output), "OpenAI"); const text = requireText(collectText(data.output_text ?? data.output), "OpenAI");
return { text, rawResponseJson: JSON.stringify(data) }; return { text, rawResponseJson: JSON.stringify(data), rateLimitState: response.rateLimitState };
}, },
}, },
anthropic: { anthropic: {
@@ -87,15 +112,15 @@ export const providerAdapters: ProviderAdapterRegistry = {
getRateLimitState: providerRateLimitUnknown, getRateLimitState: providerRateLimitUnknown,
async checkHealth(input) { async checkHealth(input) {
const profile = requireApiKeyProfile(input.authProfile, "Anthropic"); const profile = requireApiKeyProfile(input.authProfile, "Anthropic");
await getJson("https://api.anthropic.com/v1/models", { const response = await getJson("https://api.anthropic.com/v1/models", {
"x-api-key": profile.apiKey, "x-api-key": profile.apiKey,
"anthropic-version": "2023-06-01", "anthropic-version": "2023-06-01",
}); });
return healthy(); return healthy(response.rateLimitState);
}, },
async createResponse(input) { async createResponse(input) {
const profile = requireApiKeyProfile(input.authProfile, "Anthropic"); const profile = requireApiKeyProfile(input.authProfile, "Anthropic");
const data = await postJson("https://api.anthropic.com/v1/messages", { const response = await postJson("https://api.anthropic.com/v1/messages", {
"x-api-key": profile.apiKey, "x-api-key": profile.apiKey,
"anthropic-version": "2023-06-01", "anthropic-version": "2023-06-01",
}, { }, {
@@ -104,8 +129,9 @@ export const providerAdapters: ProviderAdapterRegistry = {
...(input.instructions ? { system: input.instructions } : {}), ...(input.instructions ? { system: input.instructions } : {}),
messages: [{ role: "user", content: input.prompt }], messages: [{ role: "user", content: input.prompt }],
}); });
const data = response.data;
const text = requireText(collectText(data.content), "Anthropic"); const text = requireText(collectText(data.content), "Anthropic");
return { text, rawResponseJson: JSON.stringify(data) }; return { text, rawResponseJson: JSON.stringify(data), rateLimitState: response.rateLimitState };
}, },
}, },
xai: openAiCompatibleAdapter({ xai: openAiCompatibleAdapter({
@@ -131,8 +157,8 @@ export const providerAdapters: ProviderAdapterRegistry = {
getRateLimitState: providerRateLimitUnknown, getRateLimitState: providerRateLimitUnknown,
async checkHealth(input) { async checkHealth(input) {
const profile = requireAzureProfile(input.authProfile); const profile = requireAzureProfile(input.authProfile);
await getJson(`${profile.baseUrl}/openai/deployments?api-version=2024-10-21`, { "api-key": profile.apiKey }); const response = await getJson(`${profile.baseUrl}/openai/deployments?api-version=2024-10-21`, { "api-key": profile.apiKey });
return healthy(); return healthy(response.rateLimitState);
}, },
}, },
"aws-bedrock": { "aws-bedrock": {
@@ -143,7 +169,7 @@ export const providerAdapters: ProviderAdapterRegistry = {
getRateLimitState: providerRateLimitUnknown, getRateLimitState: providerRateLimitUnknown,
checkHealth: async (input) => { checkHealth: async (input) => {
requireBedrockProfile(input.authProfile); requireBedrockProfile(input.authProfile);
return healthy(); return healthy(null);
}, },
}, },
"vertex-ai": { "vertex-ai": {
@@ -154,12 +180,12 @@ export const providerAdapters: ProviderAdapterRegistry = {
getRateLimitState: providerRateLimitUnknown, getRateLimitState: providerRateLimitUnknown,
checkHealth: async (input) => { checkHealth: async (input) => {
requireVertexProfile(input.authProfile); requireVertexProfile(input.authProfile);
return healthy(); return healthy(null);
}, },
}, },
}; };
async function postJson(endpoint: string, headers: Record<string, string>, body: JsonRecord): Promise<JsonRecord> { async function postJson(endpoint: string, headers: Record<string, string>, body: JsonRecord): Promise<ProviderJsonResponse> {
return requestJson(endpoint, { return requestJson(endpoint, {
method: "POST", method: "POST",
headers: { headers: {
@@ -170,11 +196,11 @@ async function postJson(endpoint: string, headers: Record<string, string>, body:
}); });
} }
async function getJson(endpoint: string, headers: Record<string, string>): Promise<JsonRecord> { async function getJson(endpoint: string, headers: Record<string, string>): Promise<ProviderJsonResponse> {
return requestJson(endpoint, { method: "GET", headers }); return requestJson(endpoint, { method: "GET", headers });
} }
async function requestJson(endpoint: string, init: RequestInit): Promise<JsonRecord> { async function requestJson(endpoint: string, init: RequestInit): Promise<ProviderJsonResponse> {
const controller = new AbortController(); const controller = new AbortController();
const timeout = setTimeout(() => controller.abort(), PROVIDER_REQUEST_TIMEOUT_MS); const timeout = setTimeout(() => controller.abort(), PROVIDER_REQUEST_TIMEOUT_MS);
let response: Response; let response: Response;
@@ -183,13 +209,14 @@ async function requestJson(endpoint: string, init: RequestInit): Promise<JsonRec
} finally { } finally {
clearTimeout(timeout); clearTimeout(timeout);
} }
const rateLimitState = parseProviderRateLimitHeaders(response.headers);
const raw = await response.text(); const raw = await response.text();
const data = parseProviderJson(raw, response.ok); const data = parseProviderJson(raw, response.ok);
if (!response.ok) { if (!response.ok) {
const message = readProviderError(data) ?? (raw.trim() || `Provider request failed with HTTP ${response.status}.`); const message = readProviderError(data) ?? (raw.trim() || `Provider request failed with HTTP ${response.status}.`);
throw new Error(message); throw new ProviderHttpError(message, rateLimitState);
} }
return data; return { data, rateLimitState };
} }
function parseProviderJson(raw: string, requireJson: boolean): JsonRecord { function parseProviderJson(raw: string, requireJson: boolean): JsonRecord {
@@ -254,22 +281,23 @@ function openAiCompatibleAdapter({
...(input.instructions ? [{ role: "system", content: input.instructions }] : []), ...(input.instructions ? [{ role: "system", content: input.instructions }] : []),
{ role: "user", content: input.prompt }, { role: "user", content: input.prompt },
]; ];
const data = await postJson(endpoint, { const response = await postJson(endpoint, {
authorization: `Bearer ${profile.apiKey}`, authorization: `Bearer ${profile.apiKey}`,
}, { }, {
model: input.model, model: input.model,
messages, messages,
}); });
const data = response.data;
const text = requireText(collectText(data.choices), label); const text = requireText(collectText(data.choices), label);
return { text, rawResponseJson: JSON.stringify(data) }; return { text, rawResponseJson: JSON.stringify(data), rateLimitState: response.rateLimitState };
}, },
}; };
} }
async function checkOpenAiHealth(authProfile: ProviderAuthProfile, healthUrl: string): Promise<ProviderHealthResult> { async function checkOpenAiHealth(authProfile: ProviderAuthProfile, healthUrl: string): Promise<ProviderHealthResult> {
const profile = requireApiKeyProfile(authProfile, "OpenAI-compatible provider"); const profile = requireApiKeyProfile(authProfile, "OpenAI-compatible provider");
await getJson(healthUrl, { authorization: `Bearer ${profile.apiKey}` }); const response = await getJson(healthUrl, { authorization: `Bearer ${profile.apiKey}` });
return healthy(); return healthy(response.rateLimitState);
} }
function browserTokenAdapter(providerKey: string, authType: ProviderAuthType): ProviderAdapter { function browserTokenAdapter(providerKey: string, authType: ProviderAuthType): ProviderAdapter {
@@ -282,7 +310,7 @@ function browserTokenAdapter(providerKey: string, authType: ProviderAuthType): P
checkHealth: async (input) => { checkHealth: async (input) => {
if (authType === PROVIDER_AUTH_TYPE.ACCOUNT_SUBSCRIPTION) requireAccountTokenProfile(input.authProfile); if (authType === PROVIDER_AUTH_TYPE.ACCOUNT_SUBSCRIPTION) requireAccountTokenProfile(input.authProfile);
else requireAccessTokenProfile(input.authProfile); else requireAccessTokenProfile(input.authProfile);
return healthy(); return healthy(null);
}, },
}; };
} }
@@ -399,8 +427,8 @@ function requireVertexProfile(profile: ProviderAuthProfile): Extract<ProviderAut
return profile; return profile;
} }
function healthy(): ProviderHealthResult { function healthy(rateLimitState: ProviderObservedRateLimitState | null): ProviderHealthResult {
return { checkedAt: new Date().toISOString() }; return { checkedAt: new Date().toISOString(), rateLimitState };
} }
function listStaticModels(providerKey: string): Promise<ProviderModelEntry[]> { function listStaticModels(providerKey: string): Promise<ProviderModelEntry[]> {
@@ -0,0 +1,168 @@
import type { DatabaseSync } from "node:sqlite";
import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain";
import { createId } from "../../shared/ids";
const REQUEST_LIMIT_HEADERS = [
"x-ratelimit-limit-requests",
"x-ratelimit-limit",
"anthropic-ratelimit-requests-limit",
"ratelimit-limit",
] as const;
const REMAINING_REQUEST_HEADERS = [
"x-ratelimit-remaining-requests",
"x-ratelimit-remaining",
"anthropic-ratelimit-requests-remaining",
"ratelimit-remaining",
] as const;
const RESET_HEADERS = [
"x-ratelimit-reset-requests",
"x-ratelimit-reset",
"anthropic-ratelimit-requests-reset",
"ratelimit-reset",
] as const;
const RETRY_AFTER_HEADERS = ["retry-after"] as const;
const OBSERVED_HEADERS = [
...REQUEST_LIMIT_HEADERS,
...REMAINING_REQUEST_HEADERS,
...RESET_HEADERS,
...RETRY_AFTER_HEADERS,
] as const;
export interface RecordProviderRateLimitObservationInput {
orgId: string;
providerId: string;
providerKey: string;
model: string | null;
state: ProviderObservedRateLimitState;
}
export class ProviderRateLimitRegistry {
constructor(private readonly db: DatabaseSync) {}
recordObservation(input: RecordProviderRateLimitObservationInput): void {
this.db.prepare(`
insert into provider_rate_limit_observations(
id, org_id, provider_id, provider_key, model, request_limit,
remaining_requests, reset_at, retry_after_ms, raw_state_json, observed_at
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`).run(
createId("prlo"),
input.orgId,
input.providerId,
input.providerKey,
input.model,
input.state.requestLimit,
input.state.remainingRequests,
input.state.resetAt,
input.state.retryAfterMs,
input.state.rawStateJson,
new Date().toISOString(),
);
}
}
export function parseProviderRateLimitHeaders(headers: Headers, now = new Date()): ProviderObservedRateLimitState | null {
const rawHeaders = collectObservedHeaders(headers);
const requestLimit = readIntegerHeader(headers, REQUEST_LIMIT_HEADERS);
const remainingRequests = readIntegerHeader(headers, REMAINING_REQUEST_HEADERS);
const resetAt = readResetAt(headers, RESET_HEADERS, now);
const retryAfterMs = readRetryAfterMs(headers, now);
if (
requestLimit === null
&& remainingRequests === null
&& resetAt === null
&& retryAfterMs === null
) {
return null;
}
return {
requestLimit,
remainingRequests,
resetAt,
retryAfterMs,
rawStateJson: JSON.stringify(rawHeaders),
};
}
function collectObservedHeaders(headers: Headers): Record<string, string> {
const rawHeaders: Record<string, string> = {};
for (const key of OBSERVED_HEADERS) {
const value = headers.get(key);
if (value?.trim()) rawHeaders[key] = value.trim();
}
return rawHeaders;
}
function readIntegerHeader(headers: Headers, keys: readonly string[]): number | null {
for (const key of keys) {
const value = parseInteger(headers.get(key));
if (value !== null) return value;
}
return null;
}
function readResetAt(headers: Headers, keys: readonly string[], now: Date): string | null {
for (const key of keys) {
const value = parseResetValue(headers.get(key), now);
if (value !== null) return value;
}
return null;
}
function readRetryAfterMs(headers: Headers, now: Date): number | null {
for (const key of RETRY_AFTER_HEADERS) {
const value = parseRetryAfter(headers.get(key), now);
if (value !== null) return value;
}
return null;
}
function parseInteger(value: string | null): number | null {
if (!value) return null;
const normalized = value.trim();
if (!/^\d+$/.test(normalized)) return null;
const parsed = Number(normalized);
if (!Number.isSafeInteger(parsed)) return null;
return parsed;
}
function parseResetValue(value: string | null, now: Date): string | null {
if (!value?.trim()) return null;
const normalized = value.trim();
const durationMs = parseDurationMs(normalized);
if (durationMs !== null) return new Date(now.getTime() + durationMs).toISOString();
const numeric = Number(normalized);
if (Number.isFinite(numeric) && numeric >= 0) {
if (numeric > 1_000_000_000_000) return new Date(numeric).toISOString();
if (numeric > 1_000_000_000) return new Date(numeric * 1000).toISOString();
return new Date(now.getTime() + numeric * 1000).toISOString();
}
const dateMs = Date.parse(normalized);
if (!Number.isNaN(dateMs)) return new Date(dateMs).toISOString();
return null;
}
function parseRetryAfter(value: string | null, now: Date): number | null {
if (!value?.trim()) return null;
const normalized = value.trim();
const seconds = Number(normalized);
if (Number.isFinite(seconds) && seconds >= 0) return Math.round(seconds * 1000);
const dateMs = Date.parse(normalized);
if (Number.isNaN(dateMs)) return null;
return Math.max(0, dateMs - now.getTime());
}
function parseDurationMs(value: string): number | null {
const match = /^(\d+(?:\.\d+)?)(ms|s|m|h)$/.exec(value);
if (!match) return null;
const amount = Number(match[1]);
if (!Number.isFinite(amount)) return null;
const unit = match[2];
if (unit === "ms") return Math.round(amount);
if (unit === "s") return Math.round(amount * 1000);
if (unit === "m") return Math.round(amount * 60_000);
return Math.round(amount * 3_600_000);
}
@@ -10,7 +10,9 @@ import { AuditService } from "./audit-service";
import { AuthBrokerService, type ProviderSecretVault } from "./auth-broker-service"; import { AuthBrokerService, type ProviderSecretVault } from "./auth-broker-service";
import { OperationsService } from "./operations-service"; import { OperationsService } from "./operations-service";
import { SecretVault } from "./secret-vault"; import { SecretVault } from "./secret-vault";
import { providerAdapters, type ProviderAdapterRegistry } from "./provider-adapters"; import { providerAdapters, providerRateLimitStateFromError, type ProviderAdapterRegistry } from "./provider-adapters";
import { ProviderRateLimitRegistry } from "./provider-rate-limit-registry";
import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain";
const AUTH_STATE_CONFIGURED = "configured"; const AUTH_STATE_CONFIGURED = "configured";
const AUTH_STATE_REVOKED = "revoked"; const AUTH_STATE_REVOKED = "revoked";
@@ -23,6 +25,7 @@ export class ProviderService {
private readonly audit: AuditService; private readonly audit: AuditService;
private readonly authBroker: AuthBrokerService; private readonly authBroker: AuthBrokerService;
private readonly operations: OperationsService; private readonly operations: OperationsService;
private readonly rateLimitRegistry: ProviderRateLimitRegistry;
constructor( constructor(
private readonly db: DatabaseSync, private readonly db: DatabaseSync,
@@ -32,6 +35,7 @@ export class ProviderService {
this.audit = new AuditService(db); this.audit = new AuditService(db);
this.authBroker = new AuthBrokerService(db, vault); this.authBroker = new AuthBrokerService(db, vault);
this.operations = new OperationsService(db); this.operations = new OperationsService(db);
this.rateLimitRegistry = new ProviderRateLimitRegistry(db);
} }
saveCredential(input: SaveProviderCredentialInput): void { saveCredential(input: SaveProviderCredentialInput): void {
@@ -104,7 +108,8 @@ export class ProviderService {
const credential = this.loadCredential(orgId, provider); const credential = this.loadCredential(orgId, provider);
try { try {
const profile = adapter.parseAuthProfile(credential.secret); const profile = adapter.parseAuthProfile(credential.secret);
await adapter.checkHealth({ authProfile: profile, model: provider.model }); const result = await adapter.checkHealth({ authProfile: profile, model: provider.model });
this.recordObservedRateLimit(orgId, provider, provider.model, result.rateLimitState);
this.recordProviderHealth(provider.id, orgId, HEALTH_STATE_HEALTHY, null); this.recordProviderHealth(provider.id, orgId, HEALTH_STATE_HEALTHY, null);
this.audit.record({ this.audit.record({
orgId, orgId,
@@ -117,6 +122,7 @@ export class ProviderService {
}); });
} catch (error) { } catch (error) {
const message = error instanceof Error ? error.message : String(error); const message = error instanceof Error ? error.message : String(error);
this.recordObservedRateLimit(orgId, provider, provider.model, providerRateLimitStateFromError(error));
this.recordProviderHealth(provider.id, orgId, HEALTH_STATE_FAILED, message); this.recordProviderHealth(provider.id, orgId, HEALTH_STATE_FAILED, message);
throw error; throw error;
} }
@@ -245,6 +251,7 @@ export class ProviderService {
prompt, prompt,
instructions: input.instructions?.trim() || null, instructions: input.instructions?.trim() || null,
}); });
this.recordObservedRateLimit(input.orgId, provider, provider.model, result.rateLimitState);
this.db.prepare(` this.db.prepare(`
update provider_requests update provider_requests
set state = ?, response_text = ?, raw_response_json = ?, output_char_count = ?, completed_at = ? set state = ?, response_text = ?, raw_response_json = ?, output_char_count = ?, completed_at = ?
@@ -271,6 +278,7 @@ export class ProviderService {
}); });
} catch (error) { } catch (error) {
const message = error instanceof Error ? error.message : String(error); const message = error instanceof Error ? error.message : String(error);
this.recordObservedRateLimit(input.orgId, provider, provider.model, providerRateLimitStateFromError(error));
this.db.prepare(` this.db.prepare(`
update provider_requests update provider_requests
set state = ?, error_message = ?, completed_at = ? set state = ?, error_message = ?, completed_at = ?
@@ -385,6 +393,22 @@ export class ProviderService {
updated_at = excluded.updated_at updated_at = excluded.updated_at
`).run(createId("puw"), orgId, providerId, window.start, window.end, inputCharCount, outputCharCount, now, now); `).run(createId("puw"), orgId, providerId, window.start, window.end, inputCharCount, outputCharCount, now, now);
} }
private recordObservedRateLimit(
orgId: string,
provider: Pick<ProviderRow, "id" | "provider_key">,
model: string | null,
state: ProviderObservedRateLimitState | null,
): void {
if (!state) return;
this.rateLimitRegistry.recordObservation({
orgId,
providerId: provider.id,
providerKey: provider.provider_key,
model,
state,
});
}
} }
interface ProviderRow { interface ProviderRow {
@@ -247,10 +247,16 @@ export class SdlcService {
return typeof adapter.listModels === "function" && typeof adapter.getRateLimitState === "function"; return typeof adapter.listModels === "function" && typeof adapter.getRateLimitState === "function";
}); });
const schemaVersion = count(this.db, "select coalesce(max(version), 0) as count from schema_migrations"); const schemaVersion = count(this.db, "select coalesce(max(version), 0) as count from schema_migrations");
const rateLimitRegistry = count(this.db, `
select count(*) as count
from sqlite_master
where type = 'table' and name = 'provider_rate_limit_observations'
`);
return [ return [
{ key: "provider_adapter_coverage", passed: missing.length === 0, value: `${adapterKeys.length}/${catalogKeys.length}` }, { key: "provider_adapter_coverage", passed: missing.length === 0, value: `${adapterKeys.length}/${catalogKeys.length}` },
{ key: "provider_rate_limit_registry", passed: rateLimitRegistry === 1, value: rateLimitRegistry },
{ key: "provider_runtime_contract", passed: runtimeContractReady, value: runtimeContractReady }, { key: "provider_runtime_contract", passed: runtimeContractReady, value: runtimeContractReady },
{ key: "schema_version", passed: schemaVersion >= 21, value: schemaVersion }, { key: "schema_version", passed: schemaVersion >= 22, value: schemaVersion },
]; ];
} }
@@ -50,6 +50,11 @@ export class SnapshotService {
.map((provider) => ({ ...provider, enabled: toBool(provider.enabled), backgroundEnabled: toBool(provider.backgroundEnabled) })); .map((provider) => ({ ...provider, enabled: toBool(provider.enabled), backgroundEnabled: toBool(provider.backgroundEnabled) }));
const providerAuthProfiles = list<AppSnapshot["providerAuthProfiles"][number]>(this.db, "provider_auth_profiles", "created_at desc") const providerAuthProfiles = list<AppSnapshot["providerAuthProfiles"][number]>(this.db, "provider_auth_profiles", "created_at desc")
.map((profile) => ({ ...profile, active: toBool(profile.active) })); .map((profile) => ({ ...profile, active: toBool(profile.active) }));
const providerRateLimitObservations = list<AppSnapshot["providerRateLimitObservations"][number]>(
this.db,
"provider_rate_limit_observations",
"observed_at desc",
);
const providerUsageWindows = list<ProviderUsageWindowRow>(this.db, "provider_usage_windows", "window_start desc"); const providerUsageWindows = list<ProviderUsageWindowRow>(this.db, "provider_usage_windows", "window_start desc");
const providerRateLimits = buildProviderRateLimits(providers, providerUsageWindows); const providerRateLimits = buildProviderRateLimits(providers, providerUsageWindows);
const skillBindings = list<AppSnapshot["skillBindings"][number]>(this.db, "skill_bindings", "priority asc") const skillBindings = list<AppSnapshot["skillBindings"][number]>(this.db, "skill_bindings", "priority asc")
@@ -103,6 +108,7 @@ export class SnapshotService {
providerCatalog: PROVIDER_CATALOG, providerCatalog: PROVIDER_CATALOG,
providerRequests: list(this.db, "provider_requests", "created_at desc"), providerRequests: list(this.db, "provider_requests", "created_at desc"),
providerRateLimits, providerRateLimits,
providerRateLimitObservations,
skillDefinitions: list(this.db, "skill_definitions", "created_at asc"), skillDefinitions: list(this.db, "skill_definitions", "created_at asc"),
skillRevisions: list(this.db, "skill_revisions", "published_at desc"), skillRevisions: list(this.db, "skill_revisions", "published_at desc"),
skillBindings, skillBindings,
+3 -2
View File
@@ -6,7 +6,7 @@ import type { ApprovalAssignmentRecord, ApprovalRoutingRuleRecord, HiringPolicyR
import type { ChaosDrillRecord, ChaosDrillRunRecord, FeatureFlagRecord, OperationalEventRecord, SloTargetRecord } from "./operations-domain"; import type { ChaosDrillRecord, ChaosDrillRunRecord, FeatureFlagRecord, OperationalEventRecord, SloTargetRecord } from "./operations-domain";
import type { ProviderAuthProfileRecord } from "./provider-auth-domain"; import type { ProviderAuthProfileRecord } from "./provider-auth-domain";
import type { ProviderDataPolicy, ProviderExecutionMode } from "./provider-policy-domain"; import type { ProviderDataPolicy, ProviderExecutionMode } from "./provider-policy-domain";
import type { ProviderRateLimitStateRecord } from "./provider-runtime-domain"; import type { ProviderRateLimitObservationRecord, ProviderRateLimitStateRecord } from "./provider-runtime-domain";
import type { RollbackSnapshotRecord } from "./rollback-domain"; import type { RollbackSnapshotRecord } from "./rollback-domain";
import type { EvalCaseRecord, EvalRunRecord, EvalSuiteRecord, PrBotGateRecord, SdlcReviewAgentRecord, SdlcReviewRunRecord } from "./sdlc-domain"; import type { EvalCaseRecord, EvalRunRecord, EvalSuiteRecord, PrBotGateRecord, SdlcReviewAgentRecord, SdlcReviewRunRecord } from "./sdlc-domain";
export { BOARD_INTERVENTION_ACTION, BOARD_INTERVENTION_MODE, CHAT_THREAD_BOARD_STATE, type BoardInterventionAction, type BoardInterventionMode, type BoardInterventionRecord, type ChatThreadBoardState } from "./board-domain"; export { BOARD_INTERVENTION_ACTION, BOARD_INTERVENTION_MODE, CHAT_THREAD_BOARD_STATE, type BoardInterventionAction, type BoardInterventionMode, type BoardInterventionRecord, type ChatThreadBoardState } from "./board-domain";
@@ -17,7 +17,7 @@ export { ROUTE_TO_TYPE, type ApprovalAssignmentRecord, type ApprovalRoutingRuleR
export { CHAOS_DRILL_SCENARIO, CHAOS_DRILL_STATE, FEATURE_FLAG_KEY, SLO_SIGNAL_KEY, type ChaosDrillRecord, type ChaosDrillRunRecord, type ChaosDrillScenario, type ChaosDrillState, type FeatureFlagKey, type FeatureFlagRecord, type OperationalEventRecord, type SloSignalKey, type SloTargetRecord } from "./operations-domain"; export { CHAOS_DRILL_SCENARIO, CHAOS_DRILL_STATE, FEATURE_FLAG_KEY, SLO_SIGNAL_KEY, type ChaosDrillRecord, type ChaosDrillRunRecord, type ChaosDrillScenario, type ChaosDrillState, type FeatureFlagKey, type FeatureFlagRecord, type OperationalEventRecord, type SloSignalKey, type SloTargetRecord } from "./operations-domain";
export type { ProviderAuthProfileRecord } from "./provider-auth-domain"; export type { ProviderAuthProfileRecord } from "./provider-auth-domain";
export { PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE, type ProviderDataPolicy, type ProviderExecutionMode } from "./provider-policy-domain"; export { PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE, type ProviderDataPolicy, type ProviderExecutionMode } from "./provider-policy-domain";
export type { ProviderRateLimitStateRecord } from "./provider-runtime-domain"; export type { ProviderRateLimitObservationRecord, ProviderRateLimitStateRecord } from "./provider-runtime-domain";
export type { RollbackSnapshotRecord } from "./rollback-domain"; export type { RollbackSnapshotRecord } from "./rollback-domain";
export { EVAL_SUITE_KEY, PR_BOT_GATE_KEY, SDLC_REVIEW_AGENT_KEY, SDLC_REVIEW_TARGET_TYPE, SDLC_RUN_STATE, type EvalCaseRecord, type EvalRunRecord, type EvalSuiteKey, type EvalSuiteRecord, type PrBotGateKey, type PrBotGateRecord, type SdlcReviewAgentKey, type SdlcReviewAgentRecord, type SdlcReviewRunRecord, type SdlcReviewTargetType, type SdlcRunState } from "./sdlc-domain"; export { EVAL_SUITE_KEY, PR_BOT_GATE_KEY, SDLC_REVIEW_AGENT_KEY, SDLC_REVIEW_TARGET_TYPE, SDLC_RUN_STATE, type EvalCaseRecord, type EvalRunRecord, type EvalSuiteKey, type EvalSuiteRecord, type PrBotGateKey, type PrBotGateRecord, type SdlcReviewAgentKey, type SdlcReviewAgentRecord, type SdlcReviewRunRecord, type SdlcReviewTargetType, type SdlcRunState } from "./sdlc-domain";
export { SEARCH_ENTITY_TYPES, type SearchEntityType, type SearchResultRecord } from "./search-domain"; export { SEARCH_ENTITY_TYPES, type SearchEntityType, type SearchResultRecord } from "./search-domain";
@@ -473,6 +473,7 @@ export interface AppSnapshot {
providerCatalog: ProviderCatalogEntry[]; providerCatalog: ProviderCatalogEntry[];
providerRequests: ProviderRequestRecord[]; providerRequests: ProviderRequestRecord[];
providerRateLimits: ProviderRateLimitStateRecord[]; providerRateLimits: ProviderRateLimitStateRecord[];
providerRateLimitObservations: ProviderRateLimitObservationRecord[];
skillDefinitions: SkillDefinitionRecord[]; skillDefinitions: SkillDefinitionRecord[];
skillRevisions: SkillRevisionRecord[]; skillRevisions: SkillRevisionRecord[];
skillBindings: SkillBindingRecord[]; skillBindings: SkillBindingRecord[];
@@ -12,6 +12,23 @@ export interface ProviderRateLimitStateRecord {
exhausted: boolean; exhausted: boolean;
} }
export interface ProviderObservedRateLimitState {
requestLimit: number | null;
remainingRequests: number | null;
resetAt: string | null;
retryAfterMs: number | null;
rawStateJson: string;
}
export interface ProviderRateLimitObservationRecord extends ProviderObservedRateLimitState {
id: string;
orgId: string;
providerId: string;
providerKey: string;
model: string | null;
observedAt: string;
}
export function currentProviderUsageWindow(now = new Date()): { start: string; end: string } { export function currentProviderUsageWindow(now = new Date()): { start: string; end: string } {
const start = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)); const start = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1));
const end = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1)); const end = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1));
@@ -109,7 +109,12 @@ describe("provider contracts", () => {
.get(runId) as { state: string; failed_count: number; result_json: string }; .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 }> }; 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(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); expect(result.checks.every((check) => check.passed)).toBe(true);
connection.close(); connection.close();
}); });
@@ -23,7 +23,7 @@ describe("domain services", () => {
const connection = createConnection(); const connection = createConnection();
const rows = connection.db.prepare("select version from schema_migrations order by version asc").all() as Array<{ version: number }>; 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(); connection.close();
}); });
@@ -132,6 +132,7 @@ describe("migration service", () => {
create table providers(id text primary key, org_id text, credential_id text); 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_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_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 credentials(id text primary key, org_id text);
create table skill_definitions(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); 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 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_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_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_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_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'); 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.file_access_grants).toBe(1);
expect(body.tableCounts.provider_requests).toBe(1); expect(body.tableCounts.provider_requests).toBe(1);
expect(body.tableCounts.provider_usage_windows).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.compliance_exports).toBe(1);
expect(body.tableCounts.audit_replay_runs).toBe(1); expect(body.tableCounts.audit_replay_runs).toBe(1);
expect(body.tableCounts.rollback_snapshots).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.audit_replay_runs).toEqual(["replay-1"]);
expect(body.idMaps.rollback_snapshots).toEqual(["rollback-1"]); expect(body.idMaps.rollback_snapshots).toEqual(["rollback-1"]);
expect(body.idMaps.provider_usage_windows).toEqual(["usage-1"]); expect(body.idMaps.provider_usage_windows).toEqual(["usage-1"]);
expect(body.idMaps.provider_rate_limit_observations).toEqual(["quota-1"]);
connection.close(); connection.close();
}); });
@@ -64,8 +64,8 @@ describe("operations service", () => {
openai: { openai: {
...providerAdapters.openai, ...providerAdapters.openai,
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }), parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
checkHealth: async () => ({ checkedAt: new Date().toISOString() }), checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
createResponse: async () => ({ text: "accepted", rawResponseJson: "{}" }), createResponse: async () => ({ text: "accepted", rawResponseJson: "{}", rateLimitState: null }),
}, },
}; };
const providers = new ProviderService(connection.db, adapters, testVault); 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 { DatabaseConnection } from "../../src/main/database/connection";
import { OrganizationService } from "../../src/main/services/organization-service"; import { OrganizationService } from "../../src/main/services/organization-service";
import { ProviderService } from "../../src/main/services/provider-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 { SnapshotService } from "../../src/main/services/snapshot-service";
import { PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE } from "../../src/shared/domain"; import { PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE } from "../../src/shared/domain";
import { PROVIDER_CATALOG } from "../../src/shared/provider-catalog"; import { PROVIDER_CATALOG } from "../../src/shared/provider-catalog";
@@ -144,7 +144,7 @@ describe("provider service", () => {
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }), parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
checkHealth: async ({ authProfile }) => { checkHealth: async ({ authProfile }) => {
checkedProfile = authProfile; checkedProfile = authProfile;
return { checkedAt: new Date().toISOString() }; return { checkedAt: new Date().toISOString(), rateLimitState: null };
}, },
}, },
}; };
@@ -178,10 +178,11 @@ describe("provider service", () => {
openai: { openai: {
...providerAdapters.openai, ...providerAdapters.openai,
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }), 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 }) => ({ createResponse: async ({ authProfile, model, prompt }) => ({
text: `${authProfile.type}:${model}:${prompt}`, text: `${authProfile.type}:${model}:${prompt}`,
rawResponseJson: JSON.stringify({ provider: "openai" }), rawResponseJson: JSON.stringify({ provider: "openai" }),
rateLimitState: null,
}), }),
}, },
}; };
@@ -214,6 +215,116 @@ describe("provider service", () => {
connection.close(); 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 () => { test("tracks provider usage windows and enforces monthly request limits", async () => {
const connection = createConnection(); const connection = createConnection();
const organization = new OrganizationService(connection.db); const organization = new OrganizationService(connection.db);
@@ -223,8 +334,8 @@ describe("provider service", () => {
openai: { openai: {
...providerAdapters.openai, ...providerAdapters.openai,
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }), parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
checkHealth: async () => ({ checkedAt: new Date().toISOString() }), checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
createResponse: async ({ prompt }) => ({ text: `accepted:${prompt}`, rawResponseJson: "{}" }), createResponse: async ({ prompt }) => ({ text: `accepted:${prompt}`, rawResponseJson: "{}", rateLimitState: null }),
}, },
}; };
const providers = new ProviderService(connection.db, adapters, testVault); const providers = new ProviderService(connection.db, adapters, testVault);
@@ -292,10 +403,10 @@ describe("provider service", () => {
openai: { openai: {
...providerAdapters.openai, ...providerAdapters.openai,
parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }), parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }),
checkHealth: async () => ({ checkedAt: new Date().toISOString() }), checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }),
createResponse: async () => { createResponse: async () => {
adapterCalled = true; adapterCalled = true;
return { text: "blocked", rawResponseJson: "{}" }; return { text: "blocked", rawResponseJson: "{}", rateLimitState: null };
}, },
}, },
}; };