Wait using Promise.all()

app/(playground)/p/[agentId]/beta-proto/graph/server-actions.ts

@@ -21,10 +21,29 @@ import type { AgentId } from "../types";
 import { elementsToMarkdown } from "../utils/unstructured";
 import type { Graph } from "./types";
 
+const flushMetricsAndShutdown = async (lf: Langfuse, metricReader: any) => {
+	return new Promise<void>((resolve, reject) => {
+		const timeoutId = setTimeout(() => {
+			reject(new Error("Metric flush timeout after 20 seconds"));
+		}, 20000);
+
+		Promise.all([metricReader.forceFlush(), lf.shutdownAsync()])
+			.then(() => {
+				clearTimeout(timeoutId);
+				resolve();
+			})
+			.catch((error) => {
+				clearTimeout(timeoutId);
+				reject(error);
+			});
+	});
+};
+
 type GenerateArtifactStreamParams = {
 	userPrompt: string;
 	systemPrompt?: string;
 };
+
 export async function generateArtifactStream(
 	params: GenerateArtifactStreamParams,
 ) {
@@ -78,6 +97,19 @@ export async function generateArtifactStream(
 		stream.done();
 	})();
 
+	waitUntil(
+		flushMetricsAndShutdown(lf, metricReader).catch((error) => {
+			if (error.message === "Metric flush timeout after 20 seconds") {
+				console.error("Metric flush and Langfuse shutdown timed out:", error);
+			} else {
+				console.error(
+					"Error during metric flush and Langfuse shutdown:",
+					error,
+				);
+			}
+		}),
+	);
+
 	return { object: stream.value };
 }