Wait using Promise.all()

giselles-ai · Oct 18, 2024 · 85cb71f · 85cb71f
1 parent a2f04a5
commit 85cb71f
Showing 1 changed file with 31 additions and 3 deletions.
diff --git a/app/(playground)/p/[agentId]/beta-proto/graph/server-actions.ts b/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,
 ) {
@@ -61,11 +80,20 @@ export async function generateArtifactStream(
 				});
 
 				waitUntil(
-					metricReader.forceFlush().catch((error) => {
-						console.error("Failed to flush metrics:", error);
+					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,
+							);
+						}
 					}),
 				);
-				await lf.shutdownAsync();
 			},
 		});