Merge pull request #261 from giselles-ai/add-agent-activity-recording…

…-to-node-execution Refactor recording & reporting agent time usage
giselles-ai · Dec 20, 2024 · 3e61ec0 · 3e61ec0
2 parents 2cd4320 + c3abe51
commit 3e61ec0
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 40 deletions.
diff --git a/app/(playground)/p/[agentId]/contexts/execution.tsx b/app/(playground)/p/[agentId]/contexts/execution.tsx
@@ -376,7 +376,7 @@ interface ExecutionContextType {
 		executionId: ExecutionId,
 		forceRetryStepId?: StepId,
 	) => Promise<void>;
-	recordAgentUsageAction: (
+	onFinishPerformExecutionAction: (
 		startedAt: number,
 		endedAt: number,
 		totalDurationMs: number,
@@ -411,7 +411,7 @@ interface ExecutionProviderProps {
 		executionId: ExecutionId,
 		nodeId: NodeId,
 	) => Promise<StreamableValue<TextArtifactObject, unknown>>;
-	recordAgentUsageAction: (
+	onFinishPerformExecutionAction: (
 		startedAt: number,
 		endedAt: number,
 		totalDurationMs: number,
@@ -424,7 +424,7 @@ export function ExecutionProvider({
 	putExecutionAction,
 	retryStepAction,
 	executeNodeAction,
-	recordAgentUsageAction,
+	onFinishPerformExecutionAction,
 }: ExecutionProviderProps) {
 	const { dispatch, flush, graph } = useGraph();
 	const { setTab } = usePropertiesPanel();
@@ -448,6 +448,10 @@ export function ExecutionProvider({
 			stepExecution: CompletedStepExecution,
 			artifact: TextArtifact,
 		) => void;
+		onFinishPerformExecution?: (
+			endedAt: number,
+			durationMs: number,
+		) => Promise<void>;
 	}
 	const performFlowExecution = useCallback(
 		async ({
@@ -458,6 +462,7 @@ export function ExecutionProvider({
 			executeStepCallback,
 			updateArtifactCallback,
 			onStepFinish,
+			onFinishPerformExecution,
 		}: ExecuteFlowParams) => {
 			let currentExecution = initialExecution;
 			let totalFlowDurationMs = 0;
@@ -552,16 +557,10 @@ export function ExecutionProvider({
 					},
 				},
 			});
-
-			await recordAgentUsageAction(
-				currentExecution.runStartedAt,
-				runEndedAt,
-				currentExecution.durationMs,
-			);
-
+			await onFinishPerformExecution?.(runEndedAt, totalFlowDurationMs);
 			return currentExecution;
 		},
-		[dispatch, putExecutionAction, addToast, recordAgentUsageAction],
+		[dispatch, putExecutionAction, addToast],
 	);
 
 	const executeFlow = useCallback(
@@ -611,10 +610,19 @@ export function ExecutionProvider({
 						};
 					});
 				},
+				onFinishPerformExecution: (endedAt, durationMs) =>
+					onFinishPerformExecutionAction(flowRunStartedAt, endedAt, durationMs),
 			});
 			setExecution(finalExecution);
 		},
-		[setPlaygroundMode, executeStepAction, flush, graph, performFlowExecution],
+		[
+			setPlaygroundMode,
+			executeStepAction,
+			flush,
+			graph,
+			performFlowExecution,
+			onFinishPerformExecutionAction,
+		],
 	);
 
 	const retryFlowExecution = useCallback(
@@ -670,10 +678,18 @@ export function ExecutionProvider({
 						};
 					});
 				},
+				onFinishPerformExecution: (endedAt, durationMs) =>
+					onFinishPerformExecutionAction(flowRunStartedAt, endedAt, durationMs),
 			});
 			setExecution(finalExecution);
 		},
-		[graph.executionIndexes, flush, retryStepAction, performFlowExecution],
+		[
+			graph.executionIndexes,
+			flush,
+			retryStepAction,
+			performFlowExecution,
+			onFinishPerformExecutionAction,
+		],
 	);
 
 	const executeNode = useCallback(
@@ -734,6 +750,8 @@ export function ExecutionProvider({
 						},
 					});
 				},
+				onFinishPerformExecution: (endedAt, durationMs) =>
+					onFinishPerformExecutionAction(flowRunStartedAt, endedAt, durationMs),
 			});
 			setExecution(finalExecution);
 		},
@@ -744,6 +762,7 @@ export function ExecutionProvider({
 			graph.nodes,
 			performFlowExecution,
 			dispatch,
+			onFinishPerformExecutionAction,
 		],
 	);
 
@@ -754,7 +773,7 @@ export function ExecutionProvider({
 				executeFlow,
 				retryFlowExecution,
 				executeNode,
-				recordAgentUsageAction,
+				onFinishPerformExecutionAction,
 			}}
 		>
 			{children}

diff --git a/app/(playground)/p/[agentId]/page.tsx b/app/(playground)/p/[agentId]/page.tsx
@@ -1,6 +1,7 @@
 import { getTeamMembershipByAgentId } from "@/app/(auth)/lib/get-team-membership-by-agent-id";
 import { agents, db } from "@/drizzle";
 import { developerFlag } from "@/flags";
+import { toUTCDate } from "@/lib/date";
 import {
 	ExternalServiceName,
 	VercelBlobOperation,
@@ -9,7 +10,8 @@ import {
 	withCountMeasurement,
 } from "@/lib/opentelemetry";
 import { getUser } from "@/lib/supabase";
-import { recordAgentUsage } from "@/services/agents/activities";
+import { saveAgentActivity } from "@/services/agents/activities";
+import { reportAgentTimeUsage } from "@/services/usage-based-billing/report-agent-time-usage";
 import { del, list, put } from "@vercel/blob";
 import { ReactFlowProvider } from "@xyflow/react";
 import { eq } from "drizzle-orm";
@@ -197,13 +199,22 @@ export default async function Page({
 		return await executeNode(agentId, executionId, nodeId);
 	}
 
-	async function recordAgentUsageAction(
+	async function onFinishPerformExecutionAction(
 		startedAt: number,
 		endedAt: number,
 		totalDurationMs: number,
 	) {
 		"use server";
-		return await recordAgentUsage(agentId, startedAt, endedAt, totalDurationMs);
+
+		const startedAtDateUTC = toUTCDate(new Date(startedAt));
+		const endedAtDateUTC = toUTCDate(new Date(endedAt));
+		await saveAgentActivity(
+			agentId,
+			startedAtDateUTC,
+			endedAtDateUTC,
+			totalDurationMs,
+		);
+		await reportAgentTimeUsage(endedAtDateUTC);
 	}
 
 	return (
@@ -227,8 +238,10 @@ export default async function Page({
 												executeStepAction={executeStepAction}
 												putExecutionAction={putExecutionAction}
 												retryStepAction={retryStepAction}
-												recordAgentUsageAction={recordAgentUsageAction}
 												executeNodeAction={executeNodeAction}
+												onFinishPerformExecutionAction={
+													onFinishPerformExecutionAction
+												}
 											>
 												<Playground />
 											</ExecutionProvider>

diff --git a/services/agents/activities/index.ts b/services/agents/activities/index.ts
@@ -3,6 +3,5 @@ export { AGENT_TIME_CHARGE_LIMIT_MINUTES } from "./constants";
 export { AgentActivity } from "./deprecated-agent-activity";
 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/record-agent-usage.ts b/services/agents/activities/record-agent-usage.ts