taist 0.1.4 → 0.1.6

package/lib/trace-collector.js

@@ -157,34 +157,71 @@ export class TraceCollector extends EventEmitter {
     this.traceIds.clear();
   }
 
-  async stop() {
+  /**
+   * Stop the collector gracefully.
+   * Sends shutdown signal to workers and waits for them to flush before closing.
+   * @param {number} timeout - Max time to wait for workers to disconnect (default: 2000ms)
+   */
+  async stop(timeout = 2000) {
     if (!this.started) {
       return;
     }
 
+    // Send shutdown signal to all connected workers
+    const shutdownMessage = JSON.stringify({ type: 'shutdown' }) + '\n';
+    for (const socket of this.connections) {
+      try {
+        socket.write(shutdownMessage);
+      } catch {
+        // Socket may already be closed
+      }
+    }
+
+    // Wait for connections to close gracefully, or timeout
+    const startTime = Date.now();
+    while (this.connections.size > 0 && (Date.now() - startTime) < timeout) {
+      await new Promise(resolve => setTimeout(resolve, 50));
+    }
+
     return new Promise((resolve) => {
-      // Close all active connections
+      // Gracefully close any remaining connections (allows pending data to be read)
       for (const socket of this.connections) {
-        socket.destroy();
+        try {
+          socket.end();
+        } catch {
+          try { socket.destroy(); } catch { /* ignore */ }
+        }
       }
-      this.connections.clear();
-
-      // Close server
-      this.server.close(() => {
-        this.started = false;
 
-        // Clean up socket file
-        if (process.platform !== "win32") {
+      // Give a moment for final data to arrive before destroying
+      setTimeout(() => {
+        // Force destroy any sockets that didn't close gracefully
+        for (const socket of this.connections) {
           try {
-            fs.unlinkSync(this.socketPath);
+            socket.destroy();
           } catch {
             // Ignore
           }
         }
+        this.connections.clear();
+
+        // Close server
+        this.server.close(() => {
+          this.started = false;
+
+          // Clean up socket file
+          if (process.platform !== "win32") {
+            try {
+              fs.unlinkSync(this.socketPath);
+            } catch {
+              // Ignore
+            }
+          }
 
-        this.emit("stopped");
-        resolve();
-      });
+          this.emit("stopped");
+          resolve();
+        });
+      }, 100); // 100ms grace period for data to be read
     });
   }
 

package/lib/trace-reporter.js

@@ -79,6 +79,28 @@ export class TraceReporter extends EventEmitter {
       // Don't keep the process alive just for tracing
       this.socket.unref();
 
+      // Handle incoming messages from collector (e.g., shutdown signal)
+      let dataBuffer = '';
+      this.socket.on("data", (chunk) => {
+        dataBuffer += chunk.toString();
+        const lines = dataBuffer.split('\n');
+        dataBuffer = lines.pop(); // Keep incomplete line
+
+        for (const line of lines) {
+          if (line.trim()) {
+            try {
+              const message = JSON.parse(line);
+              if (message.type === 'shutdown') {
+                logger.debug("[reporter] Received shutdown signal, flushing...");
+                this._handleShutdown();
+              }
+            } catch {
+              // Ignore parse errors
+            }
+          }
+        }
+      });
+
       this.socket.on("error", (err) => {
         this.connected = false;
         this.connecting = false;
@@ -128,6 +150,85 @@ export class TraceReporter extends EventEmitter {
     }
   }
 
+  /**
+   * Handle shutdown signal from collector.
+   * Flushes buffer and waits for data to be sent before closing.
+   */
+  _handleShutdown() {
+    if (this.closed) return;
+
+    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,
+      });
+
+      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();
+      }
+    } else {
+      this._gracefulClose();
+    }
+  }
+
+  /**
+   * Close connection gracefully, allowing pending data to be sent.
+   */
+  _gracefulClose() {
+    if (!this.socket) {
+      this.close();
+      return;
+    }
+
+    // Use socket.end() for graceful TCP close (sends FIN, allows pending data)
+    // Add a small delay to allow the collector to read the data
+    this.socket.once('close', () => {
+      logger.debug("[reporter] Socket closed gracefully");
+      this.closed = true;
+      this.connected = false;
+      this.socket = null;
+    });
+
+    // Set a timeout in case close doesn't happen
+    const closeTimeout = setTimeout(() => {
+      logger.debug("[reporter] Close timeout, forcing");
+      this.close();
+    }, 500);
+    closeTimeout.unref();
+
+    try {
+      this.socket.end();
+    } catch {
+      clearTimeout(closeTimeout);
+      this.close();
+    }
+  }
+
   /**
    * Report a single trace event
    */

package/lib/vitest-reporter.js

@@ -148,16 +148,15 @@ export class TaistReporter {
         await this.collectorReady;
       }
 
-      // Give a small delay for any final traces to arrive
-      await new Promise(resolve => setTimeout(resolve, 100));
+      // Stop the collector gracefully - this sends shutdown signal to workers
+      // and waits for them to flush their traces before closing (up to 2s timeout)
+      await this.collector.stop();
 
-      // Get collected traces
+      // Get collected traces (after workers have flushed)
       if (this.options.showTrace) {
         this.results.trace = this.collector.getTraces();
       }
 
-      // Stop the collector
-      await this.collector.stop();
       this.collector = null;
     }
 

package/package.json

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