Merge pull request #260 from giselles-ai/implement-free-plan-agent-us…

…age-limit

Check agent time is available before performFlowExecution

app/(playground)/p/[agentId]/lib/errors.ts

@@ -0,0 +1,13 @@
+export class AgentTimeNotAvailableError extends Error {
+	constructor(
+		message = "Your agent time has been depleted. Please upgrade your plan to continue using this feature.",
+	) {
+		super(message);
+		this.name = this.constructor.name;
+
+		// Ensure proper stack trace in modern environments
+		if (Error.captureStackTrace) {
+			Error.captureStackTrace(this, AgentTimeNotAvailableError);
+		}
+	}
+}

app/(playground)/p/[agentId]/lib/execution.ts

@@ -1,19 +1,21 @@
 "use server";
 
-import { db } from "@/drizzle";
+import { agents, db, subscriptions, teams } from "@/drizzle";
 
 import {
 	createLogger,
 	waitForTelemetryExport,
 	withTokenMeasurement,
 } from "@/lib/opentelemetry";
+import { isAgentTimeAvailable } from "@/services/agents/activities";
 import { anthropic } from "@ai-sdk/anthropic";
 import { google } from "@ai-sdk/google";
 import { openai } from "@ai-sdk/openai";
 import { toJsonSchema } from "@valibot/to-json-schema";
 import { type LanguageModelV1, jsonSchema, streamObject } from "ai";
 import { createStreamableValue } from "ai/rsc";
 import { MockLanguageModelV1, simulateReadableStream } from "ai/test";
+import { and, eq } from "drizzle-orm";
 import HandleBars from "handlebars";
 import Langfuse from "langfuse";
 import * as v from "valibot";
@@ -29,11 +31,11 @@ import type {
 	NodeHandle,
 	NodeHandleId,
 	NodeId,
-	Step,
 	StepId,
 	TextArtifactObject,
 	TextGenerateActionContent,
 } from "../types";
+import { AgentTimeNotAvailableError } from "./errors";
 import { textGenerationPrompt } from "./prompts";
 import { langfuseModel, toErrorWithMessage } from "./utils";
 
@@ -259,6 +261,7 @@ function resolveRequirement(
 }
 
 interface ExecutionContext {
+	agentId: AgentId;
 	executionId: ExecutionId;
 	node: Node;
 	artifacts: Artifact[];
@@ -267,6 +270,10 @@ interface ExecutionContext {
 }
 
 async function performFlowExecution(context: ExecutionContext) {
+	const canPerform = await canPerformFlowExecution(context.agentId);
+	if (!canPerform) {
+		throw new AgentTimeNotAvailableError();
+	}
 	const startTime = Date.now();
 	const lf = new Langfuse();
 	const trace = lf.trace({
@@ -397,6 +404,7 @@ export async function executeStep(
 	}
 
 	const context: ExecutionContext = {
+		agentId,
 		executionId,
 		node,
 		artifacts,
@@ -408,6 +416,7 @@ export async function executeStep(
 }
 
 export async function retryStep(
+	agentId: AgentId,
 	retryExecutionSnapshotUrl: string,
 	executionId: ExecutionId,
 	stepId: StepId,
@@ -431,6 +440,7 @@ export async function retryStep(
 	}
 
 	const context: ExecutionContext = {
+		agentId,
 		executionId,
 		node,
 		artifacts,
@@ -464,6 +474,7 @@ export async function executeNode(
 	}
 
 	const context: ExecutionContext = {
+		agentId,
 		executionId,
 		node,
 		artifacts: graph.artifacts,
@@ -473,3 +484,32 @@ export async function executeNode(
 
 	return performFlowExecution(context);
 }
+
+async function canPerformFlowExecution(agentId: AgentId) {
+	const res = await db
+		.select({
+			dbId: teams.dbId,
+			name: teams.name,
+			type: teams.type,
+			activeSubscriptionId: subscriptions.id,
+		})
+		.from(teams)
+		.innerJoin(agents, eq(agents.teamDbId, teams.dbId))
+		.leftJoin(
+			subscriptions,
+			and(
+				eq(subscriptions.teamDbId, teams.dbId),
+				eq(subscriptions.status, "active"),
+			),
+		)
+		.where(eq(agents.id, agentId));
+	if (res.length === 0) {
+		throw new Error(`Agent with id ${agentId} not found`);
+	}
+	if (res.length > 1) {
+		throw new Error(`Agent with id ${agentId} is in multiple teams`);
+	}
+
+	const team = res[0];
+	return await isAgentTimeAvailable(team);
+}

app/(playground)/p/[agentId]/page.tsx

@@ -184,6 +184,7 @@ export default async function Page({
 	) {
 		"use server";
 		return await retryStep(
+			agentId,
 			retryExecutionSnapshotUrl,
 			executionId,
 			stepId,

services/agents/activities/has-enough-agent-time-charge.ts → services/agents/activities/deprecated-has-enough-agent-time-charge.ts

@@ -2,6 +2,9 @@ import { fetchCurrentTeam, isProPlan } from "@/services/teams";
 import { calculateAgentTimeUsageMs } from "./agent-time-usage";
 import { AGENT_TIME_CHARGE_LIMIT_MINUTES } from "./constants";
 
+/**
+ * @deprecated
+ */
 export async function hasEnoughAgentTimeCharge() {
 	const curerntTeam = await fetchCurrentTeam();
 	// If team is on a pro plan, proceed

services/agents/activities/index.ts

@@ -1,7 +1,8 @@
 export { calculateAgentTimeUsageMs } from "./agent-time-usage";
 export { AGENT_TIME_CHARGE_LIMIT_MINUTES } from "./constants";
 export { AgentActivity } from "./deprecated-agent-activity";
-export { hasEnoughAgentTimeCharge } from "./has-enough-agent-time-charge";
+export { hasEnoughAgentTimeCharge } from "./deprecated-has-enough-agent-time-charge";
+export { isAgentTimeAvailable } from "./is-agent-time-available";
 export { recordAgentUsage } from "./record-agent-usage";
 export { saveAgentActivity } from "./save-agent-activity";
 export { getMonthlyBillingCycle } from "./utils";

services/agents/activities/is-agent-time-available.ts

@@ -0,0 +1,20 @@
+import { isProPlan } from "@/services/teams";
+import type { CurrentTeam } from "@/services/teams/types";
+import { calculateAgentTimeUsageMs } from "./agent-time-usage";
+import { AGENT_TIME_CHARGE_LIMIT_MINUTES } from "./constants";
+
+export async function isAgentTimeAvailable(currentTeam: CurrentTeam) {
+	// If team is on a pro plan, proceed
+	if (isProPlan(currentTeam)) {
+		return true;
+	}
+
+	// If team is on a free plan, check agent time usage
+	const timeUsageMs = await calculateAgentTimeUsageMs(currentTeam.dbId);
+	const limitsInMs = AGENT_TIME_CHARGE_LIMIT_MINUTES.FREE * 60 * 1000;
+	if (timeUsageMs >= limitsInMs) {
+		return false;
+	}
+
+	return true;
+}