import { AppError, ErrorCodes } from '@api/utils/errors'; import { countActiveOrganizationApiKeys, countOrganizationAutomationUsageSince, countOrganizationAutomationUsageSinceByKey, createOrganizationAutomationUsage, getActiveOrganizationPlan, getOrganizationAutomationCreditsUsedSince, listRecentOrganizationAutomationUsage } from '@cio/db/queries/organization'; import type { TOrganizationApiKey, TOrganizationApiKeyType, TPlan } from '@db/types'; import { AUTOMATION_TYPE, getMcpAutomationCategory, getMcpAutomationLimits, MCP_TOOL_CREDIT_COST, type TMcpToolName } from '@cio/utils/plans'; import { PLAN } from '@cio/utils/plans/constants'; export type OrganizationAutomationUsageSummary = { type: TOrganizationApiKeyType; planName: TPlan; billingPeriodStart: string; billingPeriodEnd: string; monthlyCreditsIncluded: number; monthlyCreditsUsed: number; monthlyCreditsRemaining: number; activeKeys: number; maxActiveKeys: number; rateLimits: ReturnType['rateLimits']; recentActions: Array<{ action: string; creditsConsumed: number; createdAt: string; }>; }; function getBillingPeriodStart(date: Date = new Date()) { return new Date(Date.UTC(date.getUTCFullYear(), date.getUTCMonth(), 1, 0, 0, 0, 0)); } function getBillingPeriodEnd(date: Date = new Date()) { return new Date(Date.UTC(date.getUTCFullYear(), date.getUTCMonth() + 1, 1, 0, 0, 0, 0)); } function getMinuteWindowStart(date: Date = new Date()) { return new Date(date.getTime() - 60_000); } async function getOrganizationPlanName(orgId: string): Promise { const activePlan = await getActiveOrganizationPlan(orgId); return (activePlan?.planName as TPlan | null) ?? 'BASIC'; } export async function assertOrganizationAutomationKeyCreationAllowed( organizationId: string, type: TOrganizationApiKeyType ): Promise { const planName = await getOrganizationPlanName(organizationId); if (type === AUTOMATION_TYPE.API && planName === PLAN.BASIC) { throw new AppError('Public API keys require an Early Adopter or Enterprise plan', ErrorCodes.UPGRADE_REQUIRED, 403); } if (type !== AUTOMATION_TYPE.MCP) { return; } const limits = getMcpAutomationLimits(planName); const activeKeys = await countActiveOrganizationApiKeys(organizationId, type); if (activeKeys >= limits.maxActiveKeys) { throw new AppError( `Your plan allows up to ${limits.maxActiveKeys} active MCP key${limits.maxActiveKeys === 1 ? '' : 's'}`, ErrorCodes.AUTOMATION_KEY_LIMIT_EXCEEDED, 403 ); } } export async function getOrganizationAutomationUsageSummaryService( organizationId: string, type: TOrganizationApiKeyType ): Promise { const planName = await getOrganizationPlanName(organizationId); const billingPeriodStart = getBillingPeriodStart(); const billingPeriodEnd = getBillingPeriodEnd(); if (type !== AUTOMATION_TYPE.MCP) { return { type, planName, billingPeriodStart: billingPeriodStart.toISOString(), billingPeriodEnd: billingPeriodEnd.toISOString(), monthlyCreditsIncluded: 0, monthlyCreditsUsed: 0, monthlyCreditsRemaining: 0, activeKeys: await countActiveOrganizationApiKeys(organizationId, type), maxActiveKeys: 0, rateLimits: { perKey: { readPerMinute: 0, writePerMinute: 0, publishPerMinute: 0 }, perOrg: { readPerMinute: 0, writePerMinute: 0, publishPerMinute: 0 } }, recentActions: [] }; } const limits = getMcpAutomationLimits(planName); const [monthlyCreditsUsed, activeKeys, recentActions] = await Promise.all([ getOrganizationAutomationCreditsUsedSince(organizationId, type, billingPeriodStart.toISOString()), countActiveOrganizationApiKeys(organizationId, type), listRecentOrganizationAutomationUsage(organizationId, type, 10) ]); return { type, planName, billingPeriodStart: billingPeriodStart.toISOString(), billingPeriodEnd: billingPeriodEnd.toISOString(), monthlyCreditsIncluded: limits.monthlyCredits, monthlyCreditsUsed, monthlyCreditsRemaining: Math.max(limits.monthlyCredits - monthlyCreditsUsed, 0), activeKeys, maxActiveKeys: limits.maxActiveKeys, rateLimits: limits.rateLimits, recentActions: recentActions.map((action) => ({ action: action.action, creditsConsumed: action.creditsConsumed, createdAt: action.createdAt })) }; } export async function assertMcpAutomationUsageAllowed( automationKey: TOrganizationApiKey, toolName: TMcpToolName ): Promise { const planName = await getOrganizationPlanName(automationKey.organizationId); const limits = getMcpAutomationLimits(planName); const category = getMcpAutomationCategory(toolName); const now = new Date(); const billingPeriodStart = getBillingPeriodStart(now).toISOString(); const minuteWindowStart = getMinuteWindowStart(now).toISOString(); const [creditsUsed, keyRequestsInWindow, orgRequestsInWindow] = await Promise.all([ getOrganizationAutomationCreditsUsedSince(automationKey.organizationId, automationKey.type, billingPeriodStart), countOrganizationAutomationUsageSinceByKey(automationKey.id, category, minuteWindowStart), countOrganizationAutomationUsageSince(automationKey.organizationId, automationKey.type, category, minuteWindowStart) ]); const creditsNeeded = MCP_TOOL_CREDIT_COST[toolName]; if (creditsUsed + creditsNeeded > limits.monthlyCredits) { throw new AppError( 'Monthly automation credits exhausted for this organization', ErrorCodes.AUTOMATION_CREDIT_LIMIT_EXCEEDED, 429 ); } const keyLimit = category === 'read' ? limits.rateLimits.perKey.readPerMinute : category === 'write' ? limits.rateLimits.perKey.writePerMinute : limits.rateLimits.perKey.publishPerMinute; const orgLimit = category === 'read' ? limits.rateLimits.perOrg.readPerMinute : category === 'write' ? limits.rateLimits.perOrg.writePerMinute : limits.rateLimits.perOrg.publishPerMinute; if (keyRequestsInWindow >= keyLimit || orgRequestsInWindow >= orgLimit) { throw new AppError('Automation rate limit exceeded', ErrorCodes.AUTOMATION_RATE_LIMIT_EXCEEDED, 429); } } export async function recordMcpAutomationUsage( automationKey: TOrganizationApiKey, toolName: TMcpToolName, metadata: Record = {} ): Promise { await createOrganizationAutomationUsage({ organizationId: automationKey.organizationId, organizationApiKeyId: automationKey.id, type: automationKey.type, action: toolName, category: getMcpAutomationCategory(toolName), creditsConsumed: MCP_TOOL_CREDIT_COST[toolName], metadata }); }