diff --git a/app/(playground)/p/[agentId]/lib/errors.ts b/app/(playground)/p/[agentId]/lib/errors.ts new file mode 100644 index 00000000..549f124a --- /dev/null +++ b/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); + } + } +} diff --git a/app/(playground)/p/[agentId]/lib/execution.ts b/app/(playground)/p/[agentId]/lib/execution.ts index 2d972c5c..44236482 100644 --- a/app/(playground)/p/[agentId]/lib/execution.ts +++ b/app/(playground)/p/[agentId]/lib/execution.ts @@ -1,12 +1,13 @@ "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"; @@ -14,6 +15,7 @@ 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); +} diff --git a/app/(playground)/p/[agentId]/page.tsx b/app/(playground)/p/[agentId]/page.tsx index b2e49662..da424318 100644 --- a/app/(playground)/p/[agentId]/page.tsx +++ b/app/(playground)/p/[agentId]/page.tsx @@ -184,6 +184,7 @@ export default async function Page({ ) { "use server"; return await retryStep( + agentId, retryExecutionSnapshotUrl, executionId, stepId, diff --git a/services/agents/activities/has-enough-agent-time-charge.ts b/services/agents/activities/deprecated-has-enough-agent-time-charge.ts similarity index 96% rename from services/agents/activities/has-enough-agent-time-charge.ts rename to services/agents/activities/deprecated-has-enough-agent-time-charge.ts index 4948e846..0985e18e 100644 --- a/services/agents/activities/has-enough-agent-time-charge.ts +++ b/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 diff --git a/services/agents/activities/index.ts b/services/agents/activities/index.ts index 84a48a3f..93152240 100644 --- a/services/agents/activities/index.ts +++ b/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"; diff --git a/services/agents/activities/is-agent-time-available.ts b/services/agents/activities/is-agent-time-available.ts new file mode 100644 index 00000000..09c77c23 --- /dev/null +++ b/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; +}