taist 0.2.12 → 0.2.13

package/lib/trace-reporter.js

@@ -35,6 +35,8 @@ export class TraceReporter extends EventEmitter {
     this.workerId = options.workerId || process.pid;
     this.closed = false;
     this.shuttingDown = false; // Prevents race between shutdown signal and SIGTERM
+    this.pendingWrites = 0; // Track writes that have been initiated but callback hasn't fired
+    this.pendingWriteResolvers = []; // Resolvers to call when pendingWrites reaches 0
 
     logger.debug("[reporter] Created with socketPath:", this.socketPath, "flushImmediate:", this.flushImmediate);
 
@@ -162,50 +164,84 @@ export class TraceReporter extends EventEmitter {
 
   /**
    * Handle shutdown signal from collector.
-   * Flushes buffer and waits for data to be sent before closing.
+   * Waits for pending writes and flushes buffer before closing.
    */
   _handleShutdown() {
     if (this.closed || this.shuttingDown) return;
 
+    const lifecycleDebug = process.env.TAIST_TRACE_LIFECYCLE === 'true';
+
     // Set flag immediately to prevent exit handlers from interfering
     this.shuttingDown = true;
     this._stopFlushTimer();
 
-    // Flush and wait for data to be sent
-    if (this.buffer.length > 0 && this.socket && this.connected) {
-      const traces = this.buffer.splice(0, this.buffer.length);
-      const message = JSON.stringify({
-        type: "batch",
-        workerId: this.workerId,
-        data: traces,
+    if (lifecycleDebug) {
+      console.log(`[LIFECYCLE reporter] [${Date.now()}] _handleShutdown called, pendingWrites:`, this.pendingWrites, 'buffer:', this.buffer.length);
+    }
+
+    // Wait for any pending writes to complete first
+    const waitForPendingWrites = () => {
+      return new Promise((resolve) => {
+        if (this.pendingWrites === 0) {
+          resolve();
+        } else {
+          if (lifecycleDebug) {
+            console.log('[LIFECYCLE reporter] Waiting for', this.pendingWrites, 'pending writes...');
+          }
+          this.pendingWriteResolvers.push(resolve);
+          // Safety timeout - don't wait forever
+          setTimeout(() => {
+            if (lifecycleDebug && this.pendingWrites > 0) {
+              console.log('[LIFECYCLE reporter] Timeout waiting for pending writes, still have:', this.pendingWrites);
+            }
+            resolve();
+          }, 1000);
+        }
       });
+    };
 
-      logger.debug("[reporter] Shutdown flushing", traces.length, "traces");
+    // Wait for pending writes, then flush any remaining buffer, then close
+    waitForPendingWrites().then(() => {
+      if (lifecycleDebug) {
+        console.log('[LIFECYCLE reporter] Pending writes done, flushing remaining buffer:', this.buffer.length);
+      }
 
-      try {
-        // Write data
-        const flushed = this.socket.write(message + "\n");
-
-        if (flushed) {
-          // Data was written to kernel buffer, use setImmediate to allow
-          // the event loop to process and send the data before closing
-          setImmediate(() => {
-            this._gracefulClose();
-          });
-        } else {
-          // Buffer is full, wait for drain
-          this.socket.once('drain', () => {
-            logger.debug("[reporter] Drained, closing");
-            this._gracefulClose();
-          });
+      // Flush any remaining buffer
+      if (this.buffer.length > 0 && this.socket && this.connected) {
+        const traces = this.buffer.splice(0, this.buffer.length);
+        const message = JSON.stringify({
+          type: "batch",
+          workerId: this.workerId,
+          data: traces,
+        });
+
+        logger.debug("[reporter] Shutdown flushing", traces.length, "traces");
+
+        try {
+          // Write data
+          const flushed = this.socket.write(message + "\n");
+
+          if (flushed) {
+            // Data was written to kernel buffer, use setImmediate to allow
+            // the event loop to process and send the data before closing
+            setImmediate(() => {
+              this._gracefulClose();
+            });
+          } else {
+            // Buffer is full, wait for drain
+            this.socket.once('drain', () => {
+              logger.debug("[reporter] Drained, closing");
+              this._gracefulClose();
+            });
+          }
+        } catch (err) {
+          logger.debug("[reporter] Write error:", err.message);
+          this.close();
         }
-      } catch (err) {
-        logger.debug("[reporter] Write error:", err.message);
-        this.close();
+      } else {
+        this._gracefulClose();
       }
-    } else {
-      this._gracefulClose();
-    }
+    });
   }
 
   /**
@@ -335,8 +371,31 @@ export class TraceReporter extends EventEmitter {
       data: traces,
     });
 
+    // Track pending write
+    this.pendingWrites++;
+    if (lifecycleDebug) {
+      console.log(`[LIFECYCLE reporter] [${Date.now()}] pendingWrites++, now:`, this.pendingWrites);
+    }
+
     return new Promise((resolve, reject) => {
       const writeResult = this.socket.write(message + "\n", (err) => {
+        // Write completed (success or error)
+        this.pendingWrites--;
+        if (lifecycleDebug) {
+          console.log(`[LIFECYCLE reporter] [${Date.now()}] pendingWrites--, now:`, this.pendingWrites);
+        }
+
+        // If shutdown is waiting for pending writes, notify it
+        if (this.pendingWrites === 0 && this.pendingWriteResolvers.length > 0) {
+          if (lifecycleDebug) {
+            console.log('[LIFECYCLE reporter] All pending writes complete, notifying shutdown');
+          }
+          const resolvers = this.pendingWriteResolvers.splice(0);
+          for (const resolver of resolvers) {
+            resolver();
+          }
+        }
+
         if (err) {
           if (lifecycleDebug) console.log('[LIFECYCLE reporter] WRITE CALLBACK error:', err.message);
           // Put traces back in buffer on failure

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "taist",
-  "version": "0.2.12",
+  "version": "0.2.13",
   "description": "Token-Optimized Testing Framework for AI-Assisted Development",
   "main": "index.js",
   "type": "module",