From 0388e1003b5de3ba49e75d8c65b4e5bde760cd82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sat, 23 May 2026 19:30:11 -0400 Subject: [PATCH] Add provider rate limit snapshot state --- .../src/main/services/provider-service.ts | 12 ++--- .../src/main/services/snapshot-service.ts | 35 +++++++++++++ .../src/renderer/src/pages/SettingsPage.tsx | 3 +- .../src/pages/settings/ProvidersPanel.tsx | 52 +++++++++++-------- AgentCompany-New/src/shared/domain.ts | 3 ++ .../src/shared/provider-runtime-domain.ts | 19 +++++++ .../tests/unit/provider-service.test.ts | 12 ++++- 7 files changed, 102 insertions(+), 34 deletions(-) create mode 100644 AgentCompany-New/src/shared/provider-runtime-domain.ts diff --git a/AgentCompany-New/src/main/services/provider-service.ts b/AgentCompany-New/src/main/services/provider-service.ts index d3dabff8..7149f5a9 100644 --- a/AgentCompany-New/src/main/services/provider-service.ts +++ b/AgentCompany-New/src/main/services/provider-service.ts @@ -4,6 +4,7 @@ import { BrowserWindow, session } from "electron"; import type { CreateProviderResponseInput, SaveProviderCredentialInput, UpdateProviderPolicyInput } from "../../shared/commands"; import { FEATURE_FLAG_KEY, PROVIDER_AUTH_TYPE, PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE, PROVIDER_REQUEST_STATE, SLO_SIGNAL_KEY, type ProviderAuthType, type ProviderExecutionMode } from "../../shared/domain"; import { PROVIDER_CATALOG } from "../../shared/provider-catalog"; +import { currentProviderUsageWindow } from "../../shared/provider-runtime-domain"; import { createId } from "../../shared/ids"; import { AuditService } from "./audit-service"; import { AuthBrokerService, type ProviderSecretVault } from "./auth-broker-service"; @@ -271,7 +272,7 @@ export class ProviderService { if (executionMode === PROVIDER_EXECUTION_MODE.BACKGROUND && provider.background_enabled !== 1) { throw new Error("Background execution is disabled for this provider policy."); } - const window = currentMonthlyWindow(); + const window = currentProviderUsageWindow(); const usage = this.db.prepare(` select request_count from provider_usage_windows @@ -341,7 +342,7 @@ export class ProviderService { } private recordUsage(orgId: string, providerId: string, inputCharCount: number, outputCharCount: number): void { - const window = currentMonthlyWindow(); + const window = currentProviderUsageWindow(); const now = new Date().toISOString(); this.db.prepare(` insert into provider_usage_windows( @@ -374,10 +375,3 @@ interface ProviderRow { function hashPrompt(prompt: string): string { return createHash("sha256").update(prompt).digest("hex"); } - -function currentMonthlyWindow(): { start: string; end: string } { - const now = new Date(); - const start = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)); - const end = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1)); - return { start: start.toISOString(), end: end.toISOString() }; -} diff --git a/AgentCompany-New/src/main/services/snapshot-service.ts b/AgentCompany-New/src/main/services/snapshot-service.ts index 1596948d..25b83db9 100644 --- a/AgentCompany-New/src/main/services/snapshot-service.ts +++ b/AgentCompany-New/src/main/services/snapshot-service.ts @@ -1,8 +1,13 @@ import type { DatabaseSync, StatementSync } from "node:sqlite"; import type { AppSnapshot } from "../../shared/domain"; import { PROVIDER_CATALOG } from "../../shared/provider-catalog"; +import { currentProviderUsageWindow } from "../../shared/provider-runtime-domain"; type Row = Record; +type ProviderUsageWindowRow = Pick< + AppSnapshot["providerRateLimits"][number], + "orgId" | "providerId" | "windowStart" | "windowEnd" | "requestCount" | "inputCharCount" | "outputCharCount" +>; function all(statement: StatementSync): T[] { return statement.all() as T[]; @@ -45,6 +50,8 @@ 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 providerUsageWindows = list(this.db, "provider_usage_windows", "window_start desc"); + const providerRateLimits = buildProviderRateLimits(providers, providerUsageWindows); const skillBindings = list(this.db, "skill_bindings", "priority asc") .map((binding) => ({ ...binding, enabled: toBool(binding.enabled) })); const files = list(this.db, "file_objects", "created_at desc") @@ -95,6 +102,7 @@ export class SnapshotService { providerAuthProfiles, providerCatalog: PROVIDER_CATALOG, providerRequests: list(this.db, "provider_requests", "created_at desc"), + providerRateLimits, skillDefinitions: list(this.db, "skill_definitions", "created_at asc"), skillRevisions: list(this.db, "skill_revisions", "published_at desc"), skillBindings, @@ -140,3 +148,30 @@ export class SnapshotService { return known[0] ?? null; } } + +function buildProviderRateLimits( + providers: AppSnapshot["providers"], + providerUsageWindows: ProviderUsageWindowRow[], +): AppSnapshot["providerRateLimits"] { + const window = currentProviderUsageWindow(); + return providers.map((provider) => { + const usage = providerUsageWindows.find((item) => item.providerId === provider.id && item.windowStart === window.start); + const requestCount = usage?.requestCount ?? 0; + const inputCharCount = usage?.inputCharCount ?? 0; + const outputCharCount = usage?.outputCharCount ?? 0; + const remainingRequestCount = Math.max(0, provider.monthlyRequestLimit - requestCount); + return { + orgId: provider.orgId, + providerId: provider.id, + providerKey: provider.providerKey, + windowStart: usage?.windowStart ?? window.start, + windowEnd: usage?.windowEnd ?? window.end, + monthlyRequestLimit: provider.monthlyRequestLimit, + requestCount, + inputCharCount, + outputCharCount, + remainingRequestCount, + exhausted: remainingRequestCount === 0, + }; + }); +} diff --git a/AgentCompany-New/src/renderer/src/pages/SettingsPage.tsx b/AgentCompany-New/src/renderer/src/pages/SettingsPage.tsx index 9738da04..c016a0d7 100644 --- a/AgentCompany-New/src/renderer/src/pages/SettingsPage.tsx +++ b/AgentCompany-New/src/renderer/src/pages/SettingsPage.tsx @@ -65,6 +65,7 @@ export function SettingsPage({ snapshot, actions }: PageProps) { const employees = useMemo(() => snapshot.employees.filter((item) => item.orgId === orgId), [snapshot.employees, orgId]); const agents = useMemo(() => snapshot.agents.filter((item) => item.orgId === orgId), [snapshot.agents, orgId]); const providers = useMemo(() => snapshot.providers.filter((item) => item.orgId === orgId), [snapshot.providers, orgId]); + const providerRateLimits = useMemo(() => snapshot.providerRateLimits.filter((item) => item.orgId === orgId), [snapshot.providerRateLimits, orgId]); const skills = useMemo(() => snapshot.skillDefinitions.filter((item) => item.orgId === orgId), [snapshot.skillDefinitions, orgId]); const skillRevisions = useMemo(() => { const skillIds = new Set(skills.map((skill) => skill.id)); @@ -133,7 +134,7 @@ export function SettingsPage({ snapshot, actions }: PageProps) { providers={providers} catalog={snapshot.providerCatalog} agents={agents} - providerRequests={providerRequests} + providerRateLimits={providerRateLimits} actions={actions} /> )} diff --git a/AgentCompany-New/src/renderer/src/pages/settings/ProvidersPanel.tsx b/AgentCompany-New/src/renderer/src/pages/settings/ProvidersPanel.tsx index c8b4dfd1..15c360ee 100644 --- a/AgentCompany-New/src/renderer/src/pages/settings/ProvidersPanel.tsx +++ b/AgentCompany-New/src/renderer/src/pages/settings/ProvidersPanel.tsx @@ -4,8 +4,8 @@ import { Dialog } from "../../shared/Dialog"; import type { AgentRecord, ProviderCatalogEntry, + ProviderRateLimitStateRecord, ProviderRecord, - ProviderRequestRecord, } from "../../../../shared/domain"; import { PROVIDER_DATA_POLICY } from "../../../../shared/domain"; import { ProviderLogo } from "./ProviderLogo"; @@ -16,7 +16,7 @@ interface ProvidersPanelProps { providers: ProviderRecord[]; catalog: ProviderCatalogEntry[]; agents: AgentRecord[]; - providerRequests: ProviderRequestRecord[]; + providerRateLimits: ProviderRateLimitStateRecord[]; actions: SnapshotActions; } @@ -29,7 +29,7 @@ interface CloudFields { const CLOUD_DEFAULTS: CloudFields = { baseUrl: "", region: "", projectId: "", deployment: "" }; -export function ProvidersPanel({ orgId, providers, catalog, agents, providerRequests, actions }: ProvidersPanelProps) { +export function ProvidersPanel({ orgId, providers, catalog, agents, providerRateLimits, actions }: ProvidersPanelProps) { const [active, setActive] = useState(null); const [signInLaunched, setSignInLaunched] = useState(false); const connected = providers.filter((provider) => provider.enabled); @@ -82,7 +82,7 @@ export function ProvidersPanel({ orgId, providers, catalog, agents, providerRequ @@ -122,12 +122,12 @@ export function ProvidersPanel({ orgId, providers, catalog, agents, providerRequ function ProviderPolicyTable({ orgId, providers, - providerRequests, + providerRateLimits, actions, }: { orgId: string; providers: ProviderRecord[]; - providerRequests: ProviderRequestRecord[]; + providerRateLimits: ProviderRateLimitStateRecord[]; actions: SnapshotActions; }) { return ( @@ -140,16 +140,19 @@ function ProviderPolicyTable({ >
Provider
Data policy
Background
Monthly limit
Usage
- {providers.map((provider, index) => ( - request.providerId === provider.id && inCurrentUtcMonth(request.createdAt))} - last={index === providers.length - 1} - actions={actions} - /> - ))} + {providers.map((provider, index) => { + const rateLimit = requireProviderRateLimit(provider, providerRateLimits); + return ( + + ); + })} ); @@ -158,13 +161,13 @@ function ProviderPolicyTable({ function ProviderPolicyRow({ orgId, provider, - requests, + rateLimit, last, actions, }: { orgId: string; provider: ProviderRecord; - requests: ProviderRequestRecord[]; + rateLimit: ProviderRateLimitStateRecord; last: boolean; actions: SnapshotActions; }) { @@ -228,7 +231,7 @@ function ProviderPolicyRow({ onChange={(event) => setMonthlyRequestLimit(event.target.value)} />
-
{requests.length} / {provider.monthlyRequestLimit}
+
{rateLimit.requestCount} / {rateLimit.monthlyRequestLimit}
@@ -439,8 +442,11 @@ function readableAuth(authType: ProviderRecord["authType"]): string { return authType; } -function inCurrentUtcMonth(iso: string): boolean { - const value = new Date(iso); - const now = new Date(); - return value.getUTCFullYear() === now.getUTCFullYear() && value.getUTCMonth() === now.getUTCMonth(); +function requireProviderRateLimit( + provider: ProviderRecord, + rateLimits: ProviderRateLimitStateRecord[], +): ProviderRateLimitStateRecord { + const rateLimit = rateLimits.find((item) => item.providerId === provider.id); + if (!rateLimit) throw new Error(`Provider rate limit state missing for ${provider.id}.`); + return rateLimit; } diff --git a/AgentCompany-New/src/shared/domain.ts b/AgentCompany-New/src/shared/domain.ts index de0ce057..b4b1e720 100644 --- a/AgentCompany-New/src/shared/domain.ts +++ b/AgentCompany-New/src/shared/domain.ts @@ -6,6 +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 { 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"; @@ -16,6 +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 { 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"; @@ -470,6 +472,7 @@ export interface AppSnapshot { providerAuthProfiles: ProviderAuthProfileRecord[]; providerCatalog: ProviderCatalogEntry[]; providerRequests: ProviderRequestRecord[]; + providerRateLimits: ProviderRateLimitStateRecord[]; 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 new file mode 100644 index 00000000..ca543768 --- /dev/null +++ b/AgentCompany-New/src/shared/provider-runtime-domain.ts @@ -0,0 +1,19 @@ +export interface ProviderRateLimitStateRecord { + orgId: string; + providerId: string; + providerKey: string; + windowStart: string; + windowEnd: string; + monthlyRequestLimit: number; + requestCount: number; + inputCharCount: number; + outputCharCount: number; + remainingRequestCount: number; + exhausted: boolean; +} + +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)); + return { start: start.toISOString(), end: end.toISOString() }; +} diff --git a/AgentCompany-New/tests/unit/provider-service.test.ts b/AgentCompany-New/tests/unit/provider-service.test.ts index 45efe3da..9cd85bd9 100644 --- a/AgentCompany-New/tests/unit/provider-service.test.ts +++ b/AgentCompany-New/tests/unit/provider-service.test.ts @@ -212,13 +212,23 @@ describe("provider service", () => { where org_id = ? and provider_id = ? `).get(orgId, openAiProviderId) as { request_count: number; input_char_count: number; output_char_count: number }; expect(usage).toEqual({ request_count: 1, input_char_count: 16, output_char_count: 25 }); - expect(snapshot.load().providerRequests[0]).toMatchObject({ + const state = snapshot.load(); + expect(state.providerRequests[0]).toMatchObject({ providerId: openAiProviderId, inputCharCount: 16, outputCharCount: 25, executionMode: PROVIDER_EXECUTION_MODE.FOREGROUND, policyDecision: "allowed:standard", }); + expect(state.providerRateLimits.find((item) => item.providerId === openAiProviderId)).toMatchObject({ + providerId: openAiProviderId, + monthlyRequestLimit: 1, + requestCount: 1, + inputCharCount: 16, + outputCharCount: 25, + remainingRequestCount: 0, + exhausted: true, + }); connection.close(); });