diff --git a/AgentCompany-New/src/main/database/migrations/022-provider-rate-limit-observations.ts b/AgentCompany-New/src/main/database/migrations/022-provider-rate-limit-observations.ts new file mode 100644 index 00000000..d611bbb9 --- /dev/null +++ b/AgentCompany-New/src/main/database/migrations/022-provider-rate-limit-observations.ts @@ -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')); +`; diff --git a/AgentCompany-New/src/main/database/schema.ts b/AgentCompany-New/src/main/database/schema.ts index 05d5c73d..229ab8fa 100644 --- a/AgentCompany-New/src/main/database/schema.ts +++ b/AgentCompany-New/src/main/database/schema.ts @@ -20,6 +20,7 @@ import { sdlcControlsSchema } from "./migrations/018-sdlc-controls"; import { auditReplaySchema } from "./migrations/019-audit-replay"; import { rollbackSnapshotsSchema } from "./migrations/020-rollback-snapshots"; import { providerAuthBrokerSchema } from "./migrations/021-provider-auth-broker"; +import { providerRateLimitObservationSchema } from "./migrations/022-provider-rate-limit-observations"; interface Migration { version: number; @@ -49,6 +50,7 @@ const MIGRATIONS: Migration[] = [ { version: 19, name: "audit_replay", sql: auditReplaySchema }, { version: 20, name: "rollback_snapshots", sql: rollbackSnapshotsSchema }, { 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; diff --git a/AgentCompany-New/src/main/services/migration-service.ts b/AgentCompany-New/src/main/services/migration-service.ts index 98e94ef0..5dbe05f8 100644 --- a/AgentCompany-New/src/main/services/migration-service.ts +++ b/AgentCompany-New/src/main/services/migration-service.ts @@ -50,6 +50,7 @@ const LEGACY_TABLES = [ "providers", "provider_requests", "provider_usage_windows", + "provider_rate_limit_observations", "credentials", "skill_definitions", "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_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_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: "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" }, diff --git a/AgentCompany-New/src/main/services/provider-adapters.ts b/AgentCompany-New/src/main/services/provider-adapters.ts index 30355792..866eeadc 100644 --- a/AgentCompany-New/src/main/services/provider-adapters.ts +++ b/AgentCompany-New/src/main/services/provider-adapters.ts @@ -1,5 +1,7 @@ import { PROVIDER_AUTH_TYPE, type ProviderAuthType } from "../../shared/domain"; 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 { authProfile: ProviderAuthProfile; @@ -11,6 +13,7 @@ export interface ProviderExecutionInput { export interface ProviderExecutionResult { text: string; rawResponseJson: string; + rateLimitState: ProviderObservedRateLimitState | null; } export interface ProviderHealthInput { @@ -20,6 +23,7 @@ export interface ProviderHealthInput { export interface ProviderHealthResult { checkedAt: string; + rateLimitState: ProviderObservedRateLimitState | null; } export interface ProviderContext { @@ -31,6 +35,7 @@ export interface ProviderRateLimitState { requestLimit: number | null; remainingRequests: number | null; resetAt: string | null; + retryAfterMs: number | null; rawStateJson: string; } @@ -58,6 +63,25 @@ type JsonRecord = Record; 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 = { openai: { providerKey: "openai", @@ -68,15 +92,16 @@ export const providerAdapters: ProviderAdapterRegistry = { checkHealth: (input) => checkOpenAiHealth(input.authProfile, "https://api.openai.com/v1/models"), async createResponse(input) { 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}`, }, { model: input.model, input: input.prompt, ...(input.instructions ? { instructions: input.instructions } : {}), }); + const data = response.data; 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: { @@ -87,15 +112,15 @@ export const providerAdapters: ProviderAdapterRegistry = { getRateLimitState: providerRateLimitUnknown, async checkHealth(input) { 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, "anthropic-version": "2023-06-01", }); - return healthy(); + return healthy(response.rateLimitState); }, async createResponse(input) { 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, "anthropic-version": "2023-06-01", }, { @@ -104,8 +129,9 @@ export const providerAdapters: ProviderAdapterRegistry = { ...(input.instructions ? { system: input.instructions } : {}), messages: [{ role: "user", content: input.prompt }], }); + const data = response.data; const text = requireText(collectText(data.content), "Anthropic"); - return { text, rawResponseJson: JSON.stringify(data) }; + return { text, rawResponseJson: JSON.stringify(data), rateLimitState: response.rateLimitState }; }, }, xai: openAiCompatibleAdapter({ @@ -131,8 +157,8 @@ export const providerAdapters: ProviderAdapterRegistry = { getRateLimitState: providerRateLimitUnknown, async checkHealth(input) { const profile = requireAzureProfile(input.authProfile); - await getJson(`${profile.baseUrl}/openai/deployments?api-version=2024-10-21`, { "api-key": profile.apiKey }); - return healthy(); + const response = await getJson(`${profile.baseUrl}/openai/deployments?api-version=2024-10-21`, { "api-key": profile.apiKey }); + return healthy(response.rateLimitState); }, }, "aws-bedrock": { @@ -143,7 +169,7 @@ export const providerAdapters: ProviderAdapterRegistry = { getRateLimitState: providerRateLimitUnknown, checkHealth: async (input) => { requireBedrockProfile(input.authProfile); - return healthy(); + return healthy(null); }, }, "vertex-ai": { @@ -154,12 +180,12 @@ export const providerAdapters: ProviderAdapterRegistry = { getRateLimitState: providerRateLimitUnknown, checkHealth: async (input) => { requireVertexProfile(input.authProfile); - return healthy(); + return healthy(null); }, }, }; -async function postJson(endpoint: string, headers: Record, body: JsonRecord): Promise { +async function postJson(endpoint: string, headers: Record, body: JsonRecord): Promise { return requestJson(endpoint, { method: "POST", headers: { @@ -170,11 +196,11 @@ async function postJson(endpoint: string, headers: Record, body: }); } -async function getJson(endpoint: string, headers: Record): Promise { +async function getJson(endpoint: string, headers: Record): Promise { return requestJson(endpoint, { method: "GET", headers }); } -async function requestJson(endpoint: string, init: RequestInit): Promise { +async function requestJson(endpoint: string, init: RequestInit): Promise { const controller = new AbortController(); const timeout = setTimeout(() => controller.abort(), PROVIDER_REQUEST_TIMEOUT_MS); let response: Response; @@ -183,13 +209,14 @@ async function requestJson(endpoint: string, init: RequestInit): Promise { const profile = requireApiKeyProfile(authProfile, "OpenAI-compatible provider"); - await getJson(healthUrl, { authorization: `Bearer ${profile.apiKey}` }); - return healthy(); + const response = await getJson(healthUrl, { authorization: `Bearer ${profile.apiKey}` }); + return healthy(response.rateLimitState); } function browserTokenAdapter(providerKey: string, authType: ProviderAuthType): ProviderAdapter { @@ -282,7 +310,7 @@ function browserTokenAdapter(providerKey: string, authType: ProviderAuthType): P checkHealth: async (input) => { if (authType === PROVIDER_AUTH_TYPE.ACCOUNT_SUBSCRIPTION) requireAccountTokenProfile(input.authProfile); else requireAccessTokenProfile(input.authProfile); - return healthy(); + return healthy(null); }, }; } @@ -399,8 +427,8 @@ function requireVertexProfile(profile: ProviderAuthProfile): Extract { diff --git a/AgentCompany-New/src/main/services/provider-rate-limit-registry.ts b/AgentCompany-New/src/main/services/provider-rate-limit-registry.ts new file mode 100644 index 00000000..73b9fd01 --- /dev/null +++ b/AgentCompany-New/src/main/services/provider-rate-limit-registry.ts @@ -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 { + const rawHeaders: Record = {}; + 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); +} diff --git a/AgentCompany-New/src/main/services/provider-service.ts b/AgentCompany-New/src/main/services/provider-service.ts index dd197538..278844ef 100644 --- a/AgentCompany-New/src/main/services/provider-service.ts +++ b/AgentCompany-New/src/main/services/provider-service.ts @@ -10,7 +10,9 @@ import { AuditService } from "./audit-service"; import { AuthBrokerService, type ProviderSecretVault } from "./auth-broker-service"; import { OperationsService } from "./operations-service"; 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_REVOKED = "revoked"; @@ -23,6 +25,7 @@ export class ProviderService { private readonly audit: AuditService; private readonly authBroker: AuthBrokerService; private readonly operations: OperationsService; + private readonly rateLimitRegistry: ProviderRateLimitRegistry; constructor( private readonly db: DatabaseSync, @@ -32,6 +35,7 @@ export class ProviderService { this.audit = new AuditService(db); this.authBroker = new AuthBrokerService(db, vault); this.operations = new OperationsService(db); + this.rateLimitRegistry = new ProviderRateLimitRegistry(db); } saveCredential(input: SaveProviderCredentialInput): void { @@ -104,7 +108,8 @@ export class ProviderService { const credential = this.loadCredential(orgId, provider); try { 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.audit.record({ orgId, @@ -117,6 +122,7 @@ export class ProviderService { }); } catch (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); throw error; } @@ -245,6 +251,7 @@ export class ProviderService { prompt, instructions: input.instructions?.trim() || null, }); + this.recordObservedRateLimit(input.orgId, provider, provider.model, result.rateLimitState); this.db.prepare(` update provider_requests set state = ?, response_text = ?, raw_response_json = ?, output_char_count = ?, completed_at = ? @@ -271,6 +278,7 @@ export class ProviderService { }); } catch (error) { const message = error instanceof Error ? error.message : String(error); + this.recordObservedRateLimit(input.orgId, provider, provider.model, providerRateLimitStateFromError(error)); this.db.prepare(` update provider_requests set state = ?, error_message = ?, completed_at = ? @@ -385,6 +393,22 @@ export class ProviderService { updated_at = excluded.updated_at `).run(createId("puw"), orgId, providerId, window.start, window.end, inputCharCount, outputCharCount, now, now); } + + private recordObservedRateLimit( + orgId: string, + provider: Pick, + 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 { diff --git a/AgentCompany-New/src/main/services/sdlc-service.ts b/AgentCompany-New/src/main/services/sdlc-service.ts index 13dc0cc5..fe8bc1d5 100644 --- a/AgentCompany-New/src/main/services/sdlc-service.ts +++ b/AgentCompany-New/src/main/services/sdlc-service.ts @@ -247,10 +247,16 @@ export class SdlcService { 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 rateLimitRegistry = count(this.db, ` + select count(*) as count + from sqlite_master + where type = 'table' and name = 'provider_rate_limit_observations' + `); return [ { 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: "schema_version", passed: schemaVersion >= 21, value: schemaVersion }, + { key: "schema_version", passed: schemaVersion >= 22, value: schemaVersion }, ]; } diff --git a/AgentCompany-New/src/main/services/snapshot-service.ts b/AgentCompany-New/src/main/services/snapshot-service.ts index 25b83db9..96812256 100644 --- a/AgentCompany-New/src/main/services/snapshot-service.ts +++ b/AgentCompany-New/src/main/services/snapshot-service.ts @@ -50,6 +50,11 @@ export class SnapshotService { .map((provider) => ({ ...provider, enabled: toBool(provider.enabled), backgroundEnabled: toBool(provider.backgroundEnabled) })); const providerAuthProfiles = list(this.db, "provider_auth_profiles", "created_at desc") .map((profile) => ({ ...profile, active: toBool(profile.active) })); + const providerRateLimitObservations = list( + this.db, + "provider_rate_limit_observations", + "observed_at desc", + ); const providerUsageWindows = list(this.db, "provider_usage_windows", "window_start desc"); const providerRateLimits = buildProviderRateLimits(providers, providerUsageWindows); const skillBindings = list(this.db, "skill_bindings", "priority asc") @@ -103,6 +108,7 @@ export class SnapshotService { providerCatalog: PROVIDER_CATALOG, providerRequests: list(this.db, "provider_requests", "created_at desc"), providerRateLimits, + providerRateLimitObservations, skillDefinitions: list(this.db, "skill_definitions", "created_at asc"), skillRevisions: list(this.db, "skill_revisions", "published_at desc"), skillBindings, diff --git a/AgentCompany-New/src/shared/domain.ts b/AgentCompany-New/src/shared/domain.ts index b4b1e720..23f74768 100644 --- a/AgentCompany-New/src/shared/domain.ts +++ b/AgentCompany-New/src/shared/domain.ts @@ -6,7 +6,7 @@ import type { ApprovalAssignmentRecord, ApprovalRoutingRuleRecord, HiringPolicyR import type { ChaosDrillRecord, ChaosDrillRunRecord, FeatureFlagRecord, OperationalEventRecord, SloTargetRecord } from "./operations-domain"; import type { ProviderAuthProfileRecord } from "./provider-auth-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 { 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"; @@ -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 type { ProviderAuthProfileRecord } from "./provider-auth-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 { 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"; @@ -473,6 +473,7 @@ export interface AppSnapshot { providerCatalog: ProviderCatalogEntry[]; providerRequests: ProviderRequestRecord[]; providerRateLimits: ProviderRateLimitStateRecord[]; + providerRateLimitObservations: ProviderRateLimitObservationRecord[]; skillDefinitions: SkillDefinitionRecord[]; skillRevisions: SkillRevisionRecord[]; skillBindings: SkillBindingRecord[]; diff --git a/AgentCompany-New/src/shared/provider-runtime-domain.ts b/AgentCompany-New/src/shared/provider-runtime-domain.ts index ca543768..32b2b107 100644 --- a/AgentCompany-New/src/shared/provider-runtime-domain.ts +++ b/AgentCompany-New/src/shared/provider-runtime-domain.ts @@ -12,6 +12,23 @@ export interface ProviderRateLimitStateRecord { 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 } { const start = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)); const end = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1)); diff --git a/AgentCompany-New/tests/contract/provider-contracts.test.ts b/AgentCompany-New/tests/contract/provider-contracts.test.ts index 30589632..19094737 100644 --- a/AgentCompany-New/tests/contract/provider-contracts.test.ts +++ b/AgentCompany-New/tests/contract/provider-contracts.test.ts @@ -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(); }); diff --git a/AgentCompany-New/tests/unit/domain-services.test.ts b/AgentCompany-New/tests/unit/domain-services.test.ts index b8b8776f..9ff8d5f9 100644 --- a/AgentCompany-New/tests/unit/domain-services.test.ts +++ b/AgentCompany-New/tests/unit/domain-services.test.ts @@ -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(); }); diff --git a/AgentCompany-New/tests/unit/migration-service.test.ts b/AgentCompany-New/tests/unit/migration-service.test.ts index 111872c4..3d6d4ac9 100644 --- a/AgentCompany-New/tests/unit/migration-service.test.ts +++ b/AgentCompany-New/tests/unit/migration-service.test.ts @@ -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(); }); diff --git a/AgentCompany-New/tests/unit/operations-service.test.ts b/AgentCompany-New/tests/unit/operations-service.test.ts index dcfeba2c..a7c7f6d6 100644 --- a/AgentCompany-New/tests/unit/operations-service.test.ts +++ b/AgentCompany-New/tests/unit/operations-service.test.ts @@ -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); diff --git a/AgentCompany-New/tests/unit/provider-rate-limit-registry.test.ts b/AgentCompany-New/tests/unit/provider-rate-limit-registry.test.ts new file mode 100644 index 00000000..c2b5471c --- /dev/null +++ b/AgentCompany-New/tests/unit/provider-rate-limit-registry.test.ts @@ -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(); + }); +}); diff --git a/AgentCompany-New/tests/unit/provider-service.test.ts b/AgentCompany-New/tests/unit/provider-service.test.ts index 54d86a6b..436e9bb6 100644 --- a/AgentCompany-New/tests/unit/provider-service.test.ts +++ b/AgentCompany-New/tests/unit/provider-service.test.ts @@ -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 }; }, }, };