41 lines
1.2 KiB
TypeScript
41 lines
1.2 KiB
TypeScript
|
import type { Tool } from "ai";
|
||
|
|
import type { LimitKey } from "@/lib/subscription/config";
|
||
|
|
import {
|
||
|
|
checkAndConsume,
|
||
|
|
LimitExceededError,
|
||
|
|
} from "@/lib/subscription/service";
|
||
|
|
|
||
|
|
export function withLimit<T extends Tool>(
|
||
|
|
tool: T,
|
||
|
|
kind: LimitKey,
|
||
|
|
userId: string
|
||
|
|
): T {
|
||
|
|
const originalExecute = tool.execute;
|
||
|
|
if (!originalExecute) {
|
||
|
|
return tool;
|
||
|
|
}
|
||
|
|
|
||
|
|
// biome-ignore lint/suspicious/noExplicitAny: wrapper has to be generic across all tool signatures
|
||
|
|
const wrappedExecute = (async (input: any, options: any) => {
|
||
|
|
try {
|
||
|
|
await checkAndConsume(userId, kind);
|
||
|
|
} catch (error) {
|
||
|
|
if (error instanceof LimitExceededError) {
|
||
|
|
return {
|
||
|
|
error: "limit_exceeded",
|
||
|
|
kind: error.kind,
|
||
|
|
tier: error.tier,
|
||
|
|
used: error.used,
|
||
|
|
limit: error.limit,
|
||
|
|
message: `You've hit your ${kind.replace("_", " ")} limit for today (${error.used}/${error.limit} on ${error.tier} tier). Upgrade to Pro for higher limits.`,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
throw error;
|
||
|
|
}
|
||
|
|
return originalExecute(input, options);
|
||
|
|
// biome-ignore lint/suspicious/noExplicitAny: tool execute signature varies
|
||
|
|
}) as any;
|
||
|
|
|
||
|
|
return { ...tool, execute: wrappedExecute } as T;
|
||
|
|
}
|