From 1214f3eedb5ba3f915ce560b769d7cf9f0072223 Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Wed, 18 Dec 2024 14:58:13 +0900 Subject: [PATCH 1/4] Add saving agentActivities functionality --- .../p/[agentId]/contexts/execution.tsx | 24 +++++++++++++++-- .../p/[agentId]/lib/agent-activity.ts | 27 +++++++++++++++++++ app/(playground)/p/[agentId]/page.tsx | 16 +++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) create mode 100644 app/(playground)/p/[agentId]/lib/agent-activity.ts diff --git a/app/(playground)/p/[agentId]/contexts/execution.tsx b/app/(playground)/p/[agentId]/contexts/execution.tsx index 902da17d..67c7b370 100644 --- a/app/(playground)/p/[agentId]/contexts/execution.tsx +++ b/app/(playground)/p/[agentId]/contexts/execution.tsx @@ -251,6 +251,11 @@ interface ExecutionContextType { execution: Execution | null; execute: (nodeId: NodeId) => Promise; executeFlow: (flowId: FlowId) => Promise; + saveAgentActivityAction: ( + startedAt: number, + endedAt: number, + totalDurationMs: number, + ) => Promise; } const ExecutionContext = createContext( @@ -271,6 +276,11 @@ interface ExecutionProviderProps { ) => Promise>; executeStepAction: ExecuteStepAction; putExecutionAction: (execution: Execution) => Promise<{ blobUrl: string }>; + saveAgentActivityAction: ( + startedAt: number, + endedAt: number, + totalDurationMs: number, + ) => Promise; } export function ExecutionProvider({ @@ -278,6 +288,7 @@ export function ExecutionProvider({ executeAction, executeStepAction, putExecutionAction, + saveAgentActivityAction, }: ExecutionProviderProps) { const { dispatch, flush, graph } = useGraph(); const { setTab } = usePropertiesPanel(); @@ -426,28 +437,37 @@ export function ExecutionProvider({ // Complete flow execution setExecution(currentExecution); const { blobUrl } = await putExecutionAction(currentExecution); + const flowRunEndedAt = Date.now(); dispatch({ type: "addExecutionIndex", input: { executionIndex: { executionId, blobUrl, - completedAt: Date.now(), + completedAt: flowRunEndedAt, }, }, }); + await saveAgentActivityAction( + flowRunStartedAt, + flowRunEndedAt, + totalFlowDurationMs, + ); }, [ setPlaygroundMode, graph.flows, executeStepAction, putExecutionAction, + saveAgentActivityAction, dispatch, flush, ], ); return ( - + {children} ); diff --git a/app/(playground)/p/[agentId]/lib/agent-activity.ts b/app/(playground)/p/[agentId]/lib/agent-activity.ts new file mode 100644 index 00000000..b0a81251 --- /dev/null +++ b/app/(playground)/p/[agentId]/lib/agent-activity.ts @@ -0,0 +1,27 @@ +import { agentActivities, agents, db } from "@/drizzle"; +import { toUTCDate } from "@/lib/date"; +import type { AgentId } from "@/services/agents"; +import { eq } from "drizzle-orm"; + +export async function saveAgentActivity( + agentId: AgentId, + startedAt: number, + endedAt: number, + totalDurationMs: number, +) { + const records = await db + .select({ agentDbId: agents.dbId }) + .from(agents) + .where(eq(agents.id, agentId)); + if (records.length === 0) { + throw new Error(`Agent with id ${agentId} not found`); + } + const agentDbId = records[0].agentDbId; + + await db.insert(agentActivities).values({ + agentDbId, + startedAt: toUTCDate(new Date(startedAt)), + endedAt: toUTCDate(new Date(endedAt)), + totalDurationMs: totalDurationMs.toString(), + }); +} diff --git a/app/(playground)/p/[agentId]/page.tsx b/app/(playground)/p/[agentId]/page.tsx index 69544b45..3459fabb 100644 --- a/app/(playground)/p/[agentId]/page.tsx +++ b/app/(playground)/p/[agentId]/page.tsx @@ -24,6 +24,7 @@ import { PlaygroundModeProvider } from "./contexts/playground-mode"; import { PropertiesPanelProvider } from "./contexts/properties-panel"; import { ToastProvider } from "./contexts/toast"; import { ToolbarContextProvider } from "./contexts/toolbar"; +import { saveAgentActivity } from "./lib/agent-activity"; import { executeStep } from "./lib/execution"; import { isLatestVersion, migrateGraph } from "./lib/graph"; import { buildGraphExecutionPath, buildGraphFolderPath } from "./lib/utils"; @@ -144,6 +145,20 @@ export default async function Page({ return await action(artifactId, agentId, nodeId); } + async function saveAgentActivityAction( + startedAt: number, + endedAt: number, + totalDurationMs: number, + ) { + "use server"; + return await saveAgentActivity( + agentId, + startedAt, + endedAt, + totalDurationMs, + ); + } + async function executeStepAction( flowId: FlowId, executionId: ExecutionId, @@ -201,6 +216,7 @@ export default async function Page({ executeAction={execute} executeStepAction={executeStepAction} putExecutionAction={putExecutionAction} + saveAgentActivityAction={saveAgentActivityAction} > From 961f935fcf7065432cbf8335276e5a42c0cffe3a Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Wed, 18 Dec 2024 15:05:40 +0900 Subject: [PATCH 2/4] move save-agent-activity to services/agents/activities --- app/(playground)/p/[agentId]/page.tsx | 2 +- .../activities/{types.ts => deprecated-agent-activity.ts} | 3 +++ services/agents/activities/index.ts | 3 ++- .../agents/activities/save-agent-activity.ts | 0 4 files changed, 6 insertions(+), 2 deletions(-) rename services/agents/activities/{types.ts => deprecated-agent-activity.ts} (97%) rename app/(playground)/p/[agentId]/lib/agent-activity.ts => services/agents/activities/save-agent-activity.ts (100%) diff --git a/app/(playground)/p/[agentId]/page.tsx b/app/(playground)/p/[agentId]/page.tsx index 3459fabb..16046e21 100644 --- a/app/(playground)/p/[agentId]/page.tsx +++ b/app/(playground)/p/[agentId]/page.tsx @@ -9,6 +9,7 @@ import { withCountMeasurement, } from "@/lib/opentelemetry"; import { getUser } from "@/lib/supabase"; +import { saveAgentActivity } from "@/services/agents/activities"; import { del, list, put } from "@vercel/blob"; import { ReactFlowProvider } from "@xyflow/react"; import { eq } from "drizzle-orm"; @@ -24,7 +25,6 @@ import { PlaygroundModeProvider } from "./contexts/playground-mode"; import { PropertiesPanelProvider } from "./contexts/properties-panel"; import { ToastProvider } from "./contexts/toast"; import { ToolbarContextProvider } from "./contexts/toolbar"; -import { saveAgentActivity } from "./lib/agent-activity"; import { executeStep } from "./lib/execution"; import { isLatestVersion, migrateGraph } from "./lib/graph"; import { buildGraphExecutionPath, buildGraphFolderPath } from "./lib/utils"; diff --git a/services/agents/activities/types.ts b/services/agents/activities/deprecated-agent-activity.ts similarity index 97% rename from services/agents/activities/types.ts rename to services/agents/activities/deprecated-agent-activity.ts index fb05612b..92275e43 100644 --- a/services/agents/activities/types.ts +++ b/services/agents/activities/deprecated-agent-activity.ts @@ -1,5 +1,8 @@ import type { AgentId } from "../types"; +/** + * @deprecated + */ export class AgentActivity { private actions: AgentActivityAction[] = []; public agentId: AgentId; diff --git a/services/agents/activities/index.ts b/services/agents/activities/index.ts index 8e538a9c..05906a70 100644 --- a/services/agents/activities/index.ts +++ b/services/agents/activities/index.ts @@ -1,5 +1,6 @@ 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 { AgentActivity } from "./types"; +export { saveAgentActivity } from "./save-agent-activity"; export { getMonthlyBillingCycle } from "./utils"; diff --git a/app/(playground)/p/[agentId]/lib/agent-activity.ts b/services/agents/activities/save-agent-activity.ts similarity index 100% rename from app/(playground)/p/[agentId]/lib/agent-activity.ts rename to services/agents/activities/save-agent-activity.ts From a045a7043ca6822e2e5ec98cb12d4df12151b609 Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Wed, 18 Dec 2024 15:20:11 +0900 Subject: [PATCH 3/4] Add functionality to send agent time usage to Stripe --- .../p/[agentId]/contexts/execution.tsx | 12 +++++----- app/(playground)/p/[agentId]/page.tsx | 13 ++++------- services/agents/activities/index.ts | 1 + .../agents/activities/record-agent-usage.ts | 21 ++++++++++++++++++ .../agents/activities/save-agent-activity.ts | 9 ++++---- .../report-agent-time-usage.ts | 22 +++++++++++++++++++ 6 files changed, 58 insertions(+), 20 deletions(-) create mode 100644 services/agents/activities/record-agent-usage.ts create mode 100644 services/usage-based-billing/report-agent-time-usage.ts diff --git a/app/(playground)/p/[agentId]/contexts/execution.tsx b/app/(playground)/p/[agentId]/contexts/execution.tsx index 67c7b370..ace10dd1 100644 --- a/app/(playground)/p/[agentId]/contexts/execution.tsx +++ b/app/(playground)/p/[agentId]/contexts/execution.tsx @@ -251,7 +251,7 @@ interface ExecutionContextType { execution: Execution | null; execute: (nodeId: NodeId) => Promise; executeFlow: (flowId: FlowId) => Promise; - saveAgentActivityAction: ( + recordAgentUsageAction: ( startedAt: number, endedAt: number, totalDurationMs: number, @@ -276,7 +276,7 @@ interface ExecutionProviderProps { ) => Promise>; executeStepAction: ExecuteStepAction; putExecutionAction: (execution: Execution) => Promise<{ blobUrl: string }>; - saveAgentActivityAction: ( + recordAgentUsageAction: ( startedAt: number, endedAt: number, totalDurationMs: number, @@ -288,7 +288,7 @@ export function ExecutionProvider({ executeAction, executeStepAction, putExecutionAction, - saveAgentActivityAction, + recordAgentUsageAction, }: ExecutionProviderProps) { const { dispatch, flush, graph } = useGraph(); const { setTab } = usePropertiesPanel(); @@ -448,7 +448,7 @@ export function ExecutionProvider({ }, }, }); - await saveAgentActivityAction( + await recordAgentUsageAction( flowRunStartedAt, flowRunEndedAt, totalFlowDurationMs, @@ -459,14 +459,14 @@ export function ExecutionProvider({ graph.flows, executeStepAction, putExecutionAction, - saveAgentActivityAction, + recordAgentUsageAction, dispatch, flush, ], ); return ( {children} diff --git a/app/(playground)/p/[agentId]/page.tsx b/app/(playground)/p/[agentId]/page.tsx index 16046e21..d458d85a 100644 --- a/app/(playground)/p/[agentId]/page.tsx +++ b/app/(playground)/p/[agentId]/page.tsx @@ -9,7 +9,7 @@ import { withCountMeasurement, } from "@/lib/opentelemetry"; import { getUser } from "@/lib/supabase"; -import { saveAgentActivity } from "@/services/agents/activities"; +import { recordAgentUsage } from "@/services/agents/activities"; import { del, list, put } from "@vercel/blob"; import { ReactFlowProvider } from "@xyflow/react"; import { eq } from "drizzle-orm"; @@ -145,18 +145,13 @@ export default async function Page({ return await action(artifactId, agentId, nodeId); } - async function saveAgentActivityAction( + async function recordAgentUsageAction( startedAt: number, endedAt: number, totalDurationMs: number, ) { "use server"; - return await saveAgentActivity( - agentId, - startedAt, - endedAt, - totalDurationMs, - ); + return await recordAgentUsage(agentId, startedAt, endedAt, totalDurationMs); } async function executeStepAction( @@ -216,7 +211,7 @@ export default async function Page({ executeAction={execute} executeStepAction={executeStepAction} putExecutionAction={putExecutionAction} - saveAgentActivityAction={saveAgentActivityAction} + recordAgentUsageAction={recordAgentUsageAction} > diff --git a/services/agents/activities/index.ts b/services/agents/activities/index.ts index 05906a70..84a48a3f 100644 --- a/services/agents/activities/index.ts +++ b/services/agents/activities/index.ts @@ -2,5 +2,6 @@ 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 { recordAgentUsage } from "./record-agent-usage"; export { saveAgentActivity } from "./save-agent-activity"; export { getMonthlyBillingCycle } from "./utils"; diff --git a/services/agents/activities/record-agent-usage.ts b/services/agents/activities/record-agent-usage.ts new file mode 100644 index 00000000..1fcc4246 --- /dev/null +++ b/services/agents/activities/record-agent-usage.ts @@ -0,0 +1,21 @@ +import { toUTCDate } from "@/lib/date"; +import { reportAgentTimeUsage } from "@/services/usage-based-billing/report-agent-time-usage"; +import type { AgentId } from "../types"; +import { saveAgentActivity } from "./save-agent-activity"; + +export async function recordAgentUsage( + agentId: AgentId, + startedAt: number, + endedAt: number, + totalDurationMs: number, +) { + const startedAtDateUTC = toUTCDate(new Date(startedAt)); + const endedAtDateUTC = toUTCDate(new Date(endedAt)); + await saveAgentActivity( + agentId, + startedAtDateUTC, + endedAtDateUTC, + totalDurationMs, + ); + await reportAgentTimeUsage(endedAtDateUTC); +} diff --git a/services/agents/activities/save-agent-activity.ts b/services/agents/activities/save-agent-activity.ts index b0a81251..9c02a01d 100644 --- a/services/agents/activities/save-agent-activity.ts +++ b/services/agents/activities/save-agent-activity.ts @@ -1,12 +1,11 @@ import { agentActivities, agents, db } from "@/drizzle"; -import { toUTCDate } from "@/lib/date"; import type { AgentId } from "@/services/agents"; import { eq } from "drizzle-orm"; export async function saveAgentActivity( agentId: AgentId, - startedAt: number, - endedAt: number, + startedAt: Date, + endedAt: Date, totalDurationMs: number, ) { const records = await db @@ -20,8 +19,8 @@ export async function saveAgentActivity( await db.insert(agentActivities).values({ agentDbId, - startedAt: toUTCDate(new Date(startedAt)), - endedAt: toUTCDate(new Date(endedAt)), + startedAt: startedAt, + endedAt: endedAt, totalDurationMs: totalDurationMs.toString(), }); } diff --git a/services/usage-based-billing/report-agent-time-usage.ts b/services/usage-based-billing/report-agent-time-usage.ts new file mode 100644 index 00000000..b9932f26 --- /dev/null +++ b/services/usage-based-billing/report-agent-time-usage.ts @@ -0,0 +1,22 @@ +import { db } from "@/drizzle"; +import { stripe } from "@/services/external/stripe"; +import { fetchCurrentTeam } from "@/services/teams"; +import { processUnreportedActivities } from "@/services/usage-based-billing"; +import { AgentTimeUsageDAO } from "@/services/usage-based-billing/agent-time-usage-dao"; + +export async function reportAgentTimeUsage(targetDate: Date) { + const currentTeam = await fetchCurrentTeam(); + if (currentTeam.activeSubscriptionId == null) { + return; + } + return processUnreportedActivities( + { + teamDbId: currentTeam.dbId, + targetDate: targetDate, + }, + { + dao: new AgentTimeUsageDAO(db), + stripe: stripe, + }, + ); +} From 40439245152ac17f01d012997f41473bb2658a6d Mon Sep 17 00:00:00 2001 From: Satoshi Ebisawa Date: Thu, 19 Dec 2024 09:58:57 +0900 Subject: [PATCH 4/4] Make @/services/external/stripe be lazy initilization --- services/external/stripe/config.ts | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/services/external/stripe/config.ts b/services/external/stripe/config.ts index b58ffc4a..dfaa4bff 100644 --- a/services/external/stripe/config.ts +++ b/services/external/stripe/config.ts @@ -1,6 +1,21 @@ import { Stripe } from "stripe"; -export const stripe = new Stripe(process.env.STRIPE_SECRET_KEY as string, { - // https://github.com/stripe/stripe-node#configuration - apiVersion: "2024-11-20.acacia", -}); +let stripeInstance: Stripe | null = null; + +const handler: ProxyHandler = { + get: (_target, prop: keyof Stripe | symbol) => { + if (!stripeInstance) { + const key = process.env.STRIPE_SECRET_KEY; + if (!key) { + throw new Error("STRIPE_SECRET_KEY is not configured"); + } + stripeInstance = new Stripe(key, { + // https://github.com/stripe/stripe-node#configuration + apiVersion: "2024-11-20.acacia", + }); + } + return stripeInstance[prop as keyof Stripe]; + }, +}; + +export const stripe: Stripe = new Proxy(new Stripe("dummy"), handler);