Wrap onFinish callback by waitUntil()

giselles-ai · Nov 1, 2024 · f63a88c · f63a88c
1 parent 349eb8a
commit f63a88c
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 15 deletions.
diff --git a/app/(playground)/p/[agentId]/beta-proto/artifact/server-actions.ts b/app/(playground)/p/[agentId]/beta-proto/artifact/server-actions.ts
@@ -6,6 +6,7 @@ import { createStreamableValue } from "ai/rsc";
 
 import { getUserSubscriptionId, isRoute06User } from "@/app/(auth)/lib";
 import { metrics } from "@opentelemetry/api";
+import { waitUntil } from "@vercel/functions";
 import { Langfuse } from "langfuse";
 import { schema as artifactSchema } from "../artifact/schema";
 import type { SourceIndex } from "../source/types";
@@ -62,21 +63,43 @@ ${sourcesToText(sources)}
 				isEnabled: true,
 			},
 			onFinish: async (result) => {
-				console.log("onFinish() envirenment: ", process.env.NEXT_RUNTIME);
-				const meter = metrics.getMeter("OpenAI");
-				const tokenCounter = meter.createCounter("token_consumed", {
-					description: "Number of OpenAI API tokens consumed by each request",
-				});
-				const subscriptionId = await getUserSubscriptionId();
-				const isR06User = await isRoute06User();
-				tokenCounter.add(result.usage.totalTokens, {
-					subscriptionId,
-					isR06User,
-				});
-				generation.end({
-					output: result,
-				});
-				await lf.shutdownAsync();
+				waitUntil(
+					(async () => {
+						try {
+							console.log("Start onFinish:", process.env.NEXT_RUNTIME);
+
+							const meter = metrics.getMeter("OpenAI");
+							console.log("Meter created");
+
+							const tokenCounter = meter.createCounter("token_consumed", {
+								description:
+									"Number of OpenAI API tokens consumed by each request",
+							});
+							console.log("Counter created");
+
+							const subscriptionId = await getUserSubscriptionId();
+							const isR06User = await isRoute06User();
+							console.log("Got user info");
+
+							tokenCounter.add(result.usage.totalTokens, {
+								subscriptionId,
+								isR06User,
+							});
+							console.log("Added to counter");
+
+							await new Promise((resolve) => setTimeout(resolve, 200));
+
+							generation.end({
+								output: result,
+							});
+							await lf.shutdownAsync();
+							console.log("onFinish complete");
+						} catch (error) {
+							console.error("Error in onFinish:", error);
+							throw error;
+						}
+					})(),
+				);
 			},
 		});
 

diff --git a/package.json b/package.json
@@ -42,6 +42,7 @@
 		"@trigger.dev/sdk": "3.0.0-beta.55",
 		"@vercel/blob": "0.23.4",
 		"@vercel/flags": "2.6.1",
+		"@vercel/functions": "^1.5.0",
 		"@vercel/otel": "1.10.0",
 		"@vercel/postgres": "0.9.0",
 		"@vercel/speed-insights": "1.0.12",