diff --git a/app/(playground)/p/[agentId]/beta-proto/graph/server-actions.ts b/app/(playground)/p/[agentId]/beta-proto/graph/server-actions.ts index 6d93d67c6..75862022c 100644 --- 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((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(); }, });