Merge pull request #251 from giselles-ai/implement-agent-time-record-…

…and-reporting-to-playground-v2

Implement agent time recording and reporting to playground v2

app/(playground)/p/[agentId]/contexts/execution.tsx

@@ -32,7 +32,6 @@ import type {
 	FailedStepExecution,
 	Flow,
 	FlowId,
-	GeneratedArtifact,
 	JobExecution,
 	Node,
 	NodeId,
@@ -377,6 +376,11 @@ interface ExecutionContextType {
 		executionId: ExecutionId,
 		forceRetryStepId?: StepId,
 	) => Promise<void>;
+	recordAgentUsageAction: (
+		startedAt: number,
+		endedAt: number,
+		totalDurationMs: number,
+	) => Promise<void>;
 }
 
 const ExecutionContext = createContext<ExecutionContextType | undefined>(
@@ -407,6 +411,11 @@ interface ExecutionProviderProps {
 		executionId: ExecutionId,
 		nodeId: NodeId,
 	) => Promise<StreamableValue<TextArtifactObject, unknown>>;
+	recordAgentUsageAction: (
+		startedAt: number,
+		endedAt: number,
+		totalDurationMs: number,
+	) => Promise<void>;
 }
 
 export function ExecutionProvider({
@@ -415,6 +424,7 @@ export function ExecutionProvider({
 	putExecutionAction,
 	retryStepAction,
 	executeNodeAction,
+	recordAgentUsageAction,
 }: ExecutionProviderProps) {
 	const { dispatch, flush, graph } = useGraph();
 	const { setTab } = usePropertiesPanel();
@@ -531,20 +541,27 @@ export function ExecutionProvider({
 			});
 
 			const { blobUrl } = await putExecutionAction(executionSnapshot);
+			const runEndedAt = Date.now();
 			dispatch({
 				type: "addExecutionIndex",
 				input: {
 					executionIndex: {
 						executionId: currentExecution.id,
 						blobUrl,
-						completedAt: Date.now(),
+						completedAt: runEndedAt,
 					},
 				},
 			});
 
+			await recordAgentUsageAction(
+				currentExecution.runStartedAt,
+				runEndedAt,
+				currentExecution.durationMs,
+			);
+
 			return currentExecution;
 		},
-		[dispatch, putExecutionAction, addToast],
+		[dispatch, putExecutionAction, addToast, recordAgentUsageAction],
 	);
 
 	const executeFlow = useCallback(
@@ -737,6 +754,7 @@ export function ExecutionProvider({
 				executeFlow,
 				retryFlowExecution,
 				executeNode,
+				recordAgentUsageAction,
 			}}
 		>
 			{children}

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

@@ -9,6 +9,7 @@ import {
 	withCountMeasurement,
 } from "@/lib/opentelemetry";
 import { getUser } from "@/lib/supabase";
+import { recordAgentUsage } from "@/services/agents/activities";
 import { del, list, put } from "@vercel/blob";
 import { ReactFlowProvider } from "@xyflow/react";
 import { eq } from "drizzle-orm";
@@ -30,8 +31,6 @@ import { buildGraphExecutionPath, buildGraphFolderPath } from "./lib/utils";
 import type {
 	AgentId,
 	Artifact,
-	ArtifactId,
-	Execution,
 	ExecutionId,
 	ExecutionSnapshot,
 	FlowId,
@@ -197,6 +196,15 @@ export default async function Page({
 		return await executeNode(agentId, executionId, nodeId);
 	}
 
+	async function recordAgentUsageAction(
+		startedAt: number,
+		endedAt: number,
+		totalDurationMs: number,
+	) {
+		"use server";
+		return await recordAgentUsage(agentId, startedAt, endedAt, totalDurationMs);
+	}
+
 	return (
 		<DeveloperModeProvider developerMode={developerMode}>
 			<GraphContextProvider
@@ -218,6 +226,7 @@ export default async function Page({
 												executeStepAction={executeStepAction}
 												putExecutionAction={putExecutionAction}
 												retryStepAction={retryStepAction}
+												recordAgentUsageAction={recordAgentUsageAction}
 												executeNodeAction={executeNodeAction}
 											>
 												<Playground />

services/agents/activities/types.ts → 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;

services/agents/activities/index.ts

@@ -1,5 +1,7 @@
 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 { recordAgentUsage } from "./record-agent-usage";
+export { saveAgentActivity } from "./save-agent-activity";
 export { getMonthlyBillingCycle } from "./utils";

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);
+}

services/agents/activities/save-agent-activity.ts

@@ -0,0 +1,26 @@
+import { agentActivities, agents, db } from "@/drizzle";
+import type { AgentId } from "@/services/agents";
+import { eq } from "drizzle-orm";
+
+export async function saveAgentActivity(
+	agentId: AgentId,
+	startedAt: Date,
+	endedAt: Date,
+	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: startedAt,
+		endedAt: endedAt,
+		totalDurationMs: totalDurationMs.toString(),
+	});
+}

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<Stripe> = {
+	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);

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,
+		},
+	);
+}