From 627eb2b4bd6153f6ab1f741b702006878f0ece5f Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Thu, 19 Dec 2024 14:29:13 +0900 Subject: [PATCH 1/4] Add agentId to ExecutionContext --- app/(playground)/p/[agentId]/lib/execution.ts | 6 +++++- app/(playground)/p/[agentId]/page.tsx | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/app/(playground)/p/[agentId]/lib/execution.ts b/app/(playground)/p/[agentId]/lib/execution.ts index 2d972c5c..6d4e0806 100644 --- a/app/(playground)/p/[agentId]/lib/execution.ts +++ b/app/(playground)/p/[agentId]/lib/execution.ts @@ -29,7 +29,6 @@ import type { NodeHandle, NodeHandleId, NodeId, - Step, StepId, TextArtifactObject, TextGenerateActionContent, @@ -259,6 +258,7 @@ function resolveRequirement( } interface ExecutionContext { + agentId: AgentId; executionId: ExecutionId; node: Node; artifacts: Artifact[]; @@ -397,6 +397,7 @@ export async function executeStep( } const context: ExecutionContext = { + agentId, executionId, node, artifacts, @@ -408,6 +409,7 @@ export async function executeStep( } export async function retryStep( + agentId: AgentId, retryExecutionSnapshotUrl: string, executionId: ExecutionId, stepId: StepId, @@ -431,6 +433,7 @@ export async function retryStep( } const context: ExecutionContext = { + agentId, executionId, node, artifacts, @@ -464,6 +467,7 @@ export async function executeNode( } const context: ExecutionContext = { + agentId, executionId, node, artifacts: graph.artifacts, 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, From 21df165e543bd4e06ef484d5edad3faa8df4eb50 Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Thu, 19 Dec 2024 14:30:23 +0900 Subject: [PATCH 2/4] Add isAgentTimeAvailable method --- ...eprecated-has-enough-agent-time-charge.ts} | 3 +++ services/agents/activities/index.ts | 3 ++- .../activities/is-agent-time-available.ts | 20 +++++++++++++++++++ 3 files changed, 25 insertions(+), 1 deletion(-) rename services/agents/activities/{has-enough-agent-time-charge.ts => deprecated-has-enough-agent-time-charge.ts} (96%) create mode 100644 services/agents/activities/is-agent-time-available.ts 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; +} From 0e0b097cfb9c9a18995a1b1716a0ce562c3a0a0e Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Thu, 19 Dec 2024 14:32:15 +0900 Subject: [PATCH 3/4] Add agent time availabitily check --- app/(playground)/p/[agentId]/lib/errors.ts | 13 ++++++++ app/(playground)/p/[agentId]/lib/execution.ts | 32 ++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 app/(playground)/p/[agentId]/lib/errors.ts 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 6d4e0806..55652bb6 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 { eq } from "drizzle-orm"; import HandleBars from "handlebars"; import Langfuse from "langfuse"; import * as v from "valibot"; @@ -33,6 +35,7 @@ import type { TextArtifactObject, TextGenerateActionContent, } from "../types"; +import { AgentTimeNotAvailableError } from "./errors"; import { textGenerationPrompt } from "./prompts"; import { langfuseModel, toErrorWithMessage } from "./utils"; @@ -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({ @@ -477,3 +484,26 @@ 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, eq(subscriptions.teamDbId, teams.dbId)) + .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); +} From 426118bf1373c861a776221bb4af78cecc95df19 Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Thu, 19 Dec 2024 15:31:26 +0900 Subject: [PATCH 4/4] Fix join condition --- app/(playground)/p/[agentId]/lib/execution.ts | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/app/(playground)/p/[agentId]/lib/execution.ts b/app/(playground)/p/[agentId]/lib/execution.ts index 55652bb6..44236482 100644 --- a/app/(playground)/p/[agentId]/lib/execution.ts +++ b/app/(playground)/p/[agentId]/lib/execution.ts @@ -15,7 +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 { eq } from "drizzle-orm"; +import { and, eq } from "drizzle-orm"; import HandleBars from "handlebars"; import Langfuse from "langfuse"; import * as v from "valibot"; @@ -495,7 +495,13 @@ async function canPerformFlowExecution(agentId: AgentId) { }) .from(teams) .innerJoin(agents, eq(agents.teamDbId, teams.dbId)) - .leftJoin(subscriptions, eq(subscriptions.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`);