refactor usage quota service as a policy

This commit is contained in:
Ramnique Singh 2025-08-05 15:54:28 +05:30
parent 5b6d592d09
commit 831356a155
9 changed files with 37 additions and 37 deletions

View file

@ -14,9 +14,9 @@ import { USE_BILLING } from "../lib/feature_flags";
import { WithStringId } from "../lib/types/types";
import { getEditAgentInstructionsResponse } from "../lib/copilot/copilot";
import { container } from "@/di/container";
import { IUsageQuotaPolicyService } from "@/src/application/services/usage-quota-policy.service.interface";
import { IUsageQuotaPolicy } from "@/src/application/policies/usage-quota.policy.interface";
const usageQuotaPolicyService = container.resolve<IUsageQuotaPolicyService>('usageQuotaPolicyService');
const usageQuotaPolicy = container.resolve<IUsageQuotaPolicy>('usageQuotaPolicy');
export async function getCopilotResponseStream(
projectId: string,
@ -28,7 +28,7 @@ export async function getCopilotResponseStream(
streamId: string;
} | { billingError: string }> {
await projectAuthCheck(projectId);
await usageQuotaPolicyService.assertAndConsume(projectId);
await usageQuotaPolicy.assertAndConsume(projectId);
// Check billing authorization
const authResponse = await authorizeUserAction({
@ -39,7 +39,7 @@ export async function getCopilotResponseStream(
return { billingError: authResponse.error || 'Billing error' };
}
await usageQuotaPolicyService.assertAndConsume(projectId);
await usageQuotaPolicy.assertAndConsume(projectId);
// prepare request
const request: z.infer<typeof CopilotAPIRequest> = {
@ -71,7 +71,7 @@ export async function getCopilotAgentInstructions(
agentName: string,
): Promise<string | { billingError: string }> {
await projectAuthCheck(projectId);
await usageQuotaPolicyService.assertAndConsume(projectId);
await usageQuotaPolicy.assertAndConsume(projectId);
// Check billing authorization
const authResponse = await authorizeUserAction({

View file

@ -9,7 +9,7 @@ import { authorize, getCustomerIdForProject, logUsage } from "@/app/lib/billing"
import { USE_BILLING } from "@/app/lib/feature_flags";
import { getResponse } from "@/app/lib/agents";
import { Message, AssistantMessage, AssistantMessageWithToolCalls, ToolMessage } from "@/app/lib/types/types";
import { IUsageQuotaPolicyService } from "@/src/application/services/usage-quota-policy.service.interface";
import { IUsageQuotaPolicy } from "@/src/application/policies/usage-quota.policy.interface";
import { container } from "@/di/container";
function convert(messages: z.infer<typeof apiV1.ChatMessage>[]): z.infer<typeof Message>[] {
@ -125,8 +125,8 @@ export async function POST(
}
// assert and consume quota
const usageQuotaPolicyService = container.resolve<IUsageQuotaPolicyService>('usageQuotaPolicyService');
await usageQuotaPolicyService.assertAndConsume(session.projectId);
const usageQuotaPolicy = container.resolve<IUsageQuotaPolicy>('usageQuotaPolicy');
await usageQuotaPolicy.assertAndConsume(session.projectId);
// parse and validate the request body
let body;

View file

@ -9,7 +9,7 @@ import { CreateCachedTurnUseCase } from "@/src/application/use-cases/conversatio
import { FetchCachedTurnUseCase } from "@/src/application/use-cases/conversations/fetch-cached-turn.use-case";
import { CreateCachedTurnController } from "@/src/interface-adapters/controllers/conversations/create-cached-turn.controller";
import { RunTurnController } from "@/src/interface-adapters/controllers/conversations/run-turn.controller";
import { RedisUsageQuotaPolicyService } from "@/src/infrastructure/services/redis.usage-quota-policy.service";
import { RedisUsageQuotaPolicy } from "@/src/infrastructure/policies/redis.usage-quota.policy";
export const container = createContainer({
injectionMode: InjectionMode.PROXY,
@ -20,7 +20,7 @@ container.register({
// services
// ---
cacheService: asClass(RedisCacheService).singleton(),
usageQuotaPolicyService: asClass(RedisUsageQuotaPolicyService).singleton(),
usageQuotaPolicy: asClass(RedisUsageQuotaPolicy).singleton(),
// conversations
// ---

View file

@ -1,4 +1,4 @@
export interface IUsageQuotaPolicyService {
export interface IUsageQuotaPolicy {
// this method will throw a QuotaExceededError if the quota is exceeded
assertAndConsume(projectId: string): Promise<void>;
}

View file

@ -5,7 +5,7 @@ import { z } from "zod";
import { nanoid } from 'nanoid';
import { ICacheService } from '@/src/application/services/cache.service.interface';
import { CachedTurnRequest, Turn } from '@/src/entities/models/turn';
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
const inputSchema = z.object({
caller: z.enum(["user", "api"]),
@ -22,20 +22,20 @@ export interface ICreateCachedTurnUseCase {
export class CreateCachedTurnUseCase implements ICreateCachedTurnUseCase {
private readonly cacheService: ICacheService;
private readonly conversationsRepository: IConversationsRepository;
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
constructor({
cacheService,
conversationsRepository,
usageQuotaPolicyService,
usageQuotaPolicy,
}: {
cacheService: ICacheService,
conversationsRepository: IConversationsRepository,
usageQuotaPolicyService: IUsageQuotaPolicyService,
usageQuotaPolicy: IUsageQuotaPolicy,
}) {
this.cacheService = cacheService;
this.conversationsRepository = conversationsRepository;
this.usageQuotaPolicyService = usageQuotaPolicyService;
this.usageQuotaPolicy = usageQuotaPolicy;
}
async execute(data: z.infer<typeof inputSchema>): Promise<{ key: string }> {
@ -49,7 +49,7 @@ export class CreateCachedTurnUseCase implements ICreateCachedTurnUseCase {
const { projectId } = conversation;
// assert and consume quota
await this.usageQuotaPolicyService.assertAndConsume(projectId);
await this.usageQuotaPolicy.assertAndConsume(projectId);
// if caller is a user, ensure they are a member of project
if (data.caller === "user") {

View file

@ -4,7 +4,7 @@ import { IConversationsRepository } from "@/src/application/repositories/convers
import { z } from "zod";
import { Conversation } from "@/src/entities/models/conversation";
import { Workflow } from "@/app/lib/types/workflow_types";
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
const inputSchema = z.object({
caller: z.enum(["user", "api"]),
@ -21,17 +21,17 @@ export interface ICreateConversationUseCase {
export class CreateConversationUseCase implements ICreateConversationUseCase {
private readonly conversationsRepository: IConversationsRepository;
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
constructor({
conversationsRepository,
usageQuotaPolicyService,
usageQuotaPolicy,
}: {
conversationsRepository: IConversationsRepository,
usageQuotaPolicyService: IUsageQuotaPolicyService,
usageQuotaPolicy: IUsageQuotaPolicy,
}) {
this.conversationsRepository = conversationsRepository;
this.usageQuotaPolicyService = usageQuotaPolicyService;
this.usageQuotaPolicy = usageQuotaPolicy;
}
async execute(data: z.infer<typeof inputSchema>): Promise<z.infer<typeof Conversation>> {
@ -40,7 +40,7 @@ export class CreateConversationUseCase implements ICreateConversationUseCase {
let workflow = data.workflow;
// assert and consume quota
await this.usageQuotaPolicyService.assertAndConsume(projectId);
await this.usageQuotaPolicy.assertAndConsume(projectId);
// if caller is a user, ensure they are a member of project
if (caller === "user") {

View file

@ -4,7 +4,7 @@ import { IConversationsRepository } from "@/src/application/repositories/convers
import { z } from "zod";
import { ICacheService } from '@/src/application/services/cache.service.interface';
import { CachedTurnRequest, Turn } from '@/src/entities/models/turn';
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
const inputSchema = z.object({
caller: z.enum(["user", "api"]),
@ -20,20 +20,20 @@ export interface IFetchCachedTurnUseCase {
export class FetchCachedTurnUseCase implements IFetchCachedTurnUseCase {
private readonly cacheService: ICacheService;
private readonly conversationsRepository: IConversationsRepository;
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
constructor({
cacheService,
conversationsRepository,
usageQuotaPolicyService,
usageQuotaPolicy,
}: {
cacheService: ICacheService,
conversationsRepository: IConversationsRepository,
usageQuotaPolicyService: IUsageQuotaPolicyService,
usageQuotaPolicy: IUsageQuotaPolicy,
}) {
this.cacheService = cacheService;
this.conversationsRepository = conversationsRepository;
this.usageQuotaPolicyService = usageQuotaPolicyService;
this.usageQuotaPolicy = usageQuotaPolicy;
}
async execute(data: z.infer<typeof inputSchema>): Promise<z.infer<typeof CachedTurnRequest>> {
@ -56,7 +56,7 @@ export class FetchCachedTurnUseCase implements IFetchCachedTurnUseCase {
const { projectId } = conversation;
// assert and consume quota
await this.usageQuotaPolicyService.assertAndConsume(projectId);
await this.usageQuotaPolicy.assertAndConsume(projectId);
// if caller is a user, ensure they are a member of project
if (data.caller === "user") {

View file

@ -7,7 +7,7 @@ import { IConversationsRepository } from "@/src/application/repositories/convers
import { streamResponse } from "@/app/lib/agents";
import { z } from "zod";
import { Message } from "@/app/lib/types/types";
import { IUsageQuotaPolicyService } from '../../services/usage-quota-policy.service.interface';
import { IUsageQuotaPolicy } from '../../policies/usage-quota.policy.interface';
const inputSchema = z.object({
caller: z.enum(["user", "api"]),
@ -24,17 +24,17 @@ export interface IRunConversationTurnUseCase {
export class RunConversationTurnUseCase implements IRunConversationTurnUseCase {
private readonly conversationsRepository: IConversationsRepository;
private readonly usageQuotaPolicyService: IUsageQuotaPolicyService;
private readonly usageQuotaPolicy: IUsageQuotaPolicy;
constructor({
conversationsRepository,
usageQuotaPolicyService,
usageQuotaPolicy,
}: {
conversationsRepository: IConversationsRepository,
usageQuotaPolicyService: IUsageQuotaPolicyService,
usageQuotaPolicy: IUsageQuotaPolicy,
}) {
this.conversationsRepository = conversationsRepository;
this.usageQuotaPolicyService = usageQuotaPolicyService;
this.usageQuotaPolicy = usageQuotaPolicy;
}
async *execute(data: z.infer<typeof inputSchema>): AsyncGenerator<z.infer<typeof TurnEvent>, void, unknown> {
@ -48,7 +48,7 @@ export class RunConversationTurnUseCase implements IRunConversationTurnUseCase {
const { id: conversationId, projectId } = conversation;
// assert and consume quota
await this.usageQuotaPolicyService.assertAndConsume(projectId);
await this.usageQuotaPolicy.assertAndConsume(projectId);
// if caller is a user, ensure they are a member of project
if (data.caller === "user") {

View file

@ -1,10 +1,10 @@
import { IUsageQuotaPolicyService } from "@/src/application/services/usage-quota-policy.service.interface";
import { IUsageQuotaPolicy } from "@/src/application/policies/usage-quota.policy.interface";
import { redisClient } from "@/app/lib/redis";
import { QuotaExceededError } from "@/src/entities/errors/common";
const MAX_QUERIES_PER_MINUTE = Number(process.env.MAX_QUERIES_PER_MINUTE) || 0;
export class RedisUsageQuotaPolicyService implements IUsageQuotaPolicyService {
export class RedisUsageQuotaPolicy implements IUsageQuotaPolicy {
async assertAndConsume(projectId: string): Promise<void> {
if (MAX_QUERIES_PER_MINUTE === 0) {
return;