fix(realtime): waitUntil now correctly waits for metadata.streams to finish (#2399)

Author: ericallam

.changeset/healthy-oranges-refuse.md

@@ -0,0 +1,5 @@
+---
+"trigger.dev": patch
+---
+
+fix: waitUntil now correctly waits for metadata.streams to finish

packages/cli-v3/src/entryPoints/dev-run-worker.ts

@@ -141,11 +141,6 @@ const runMetadataManager = new StandardMetadataManager(
 runMetadata.setGlobalManager(runMetadataManager);
 const waitUntilManager = new StandardWaitUntilManager();
 waitUntil.setGlobalManager(waitUntilManager);
-// Wait for all streams to finish before completing the run
-waitUntil.register({
-  requiresResolving: () => runMetadataManager.hasActiveStreams(),
-  promise: () => runMetadataManager.waitForAllStreams(),
-});
 
 const triggerLogLevel = getEnvVar("TRIGGER_LOG_LEVEL");
 const showInternalLogs = getEnvVar("RUN_WORKER_SHOW_LOGS") === "true";
@@ -309,6 +304,12 @@ function resetExecutionEnvironment() {
   taskContext.disable();
   standardTraceContextManager.reset();
 
+  // Wait for all streams to finish before completing the run
+  waitUntil.register({
+    requiresResolving: () => runMetadataManager.hasActiveStreams(),
+    promise: () => runMetadataManager.waitForAllStreams(),
+  });
+
   log(`[${new Date().toISOString()}] Reset execution environment`);
 }
 

packages/cli-v3/src/entryPoints/managed-run-worker.ts

@@ -133,11 +133,6 @@ runMetadata.setGlobalManager(runMetadataManager);
 
 const waitUntilManager = new StandardWaitUntilManager();
 waitUntil.setGlobalManager(waitUntilManager);
-// Wait for all streams to finish before completing the run
-waitUntil.register({
-  requiresResolving: () => runMetadataManager.hasActiveStreams(),
-  promise: () => runMetadataManager.waitForAllStreams(),
-});
 
 const triggerLogLevel = getEnvVar("TRIGGER_LOG_LEVEL");
 
@@ -295,6 +290,12 @@ function resetExecutionEnvironment() {
   taskContext.disable();
   standardTraceContextManager.reset();
 
+  // Wait for all streams to finish before completing the run
+  waitUntil.register({
+    requiresResolving: () => runMetadataManager.hasActiveStreams(),
+    promise: () => runMetadataManager.waitForAllStreams(),
+  });
+
   console.log(`[${new Date().toISOString()}] Reset execution environment`);
 }