taist 0.1.4 → 0.1.5

package/lib/trace-collector.js

@@ -157,13 +157,34 @@ 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
+      // Force close any remaining connections
       for (const socket of this.connections) {
         socket.destroy();
       }

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,35 @@ export class TraceReporter extends EventEmitter {
     }
   }
 
+  /**
+   * Handle shutdown signal from collector.
+   * Immediately flushes buffer and closes connection.
+   */
+  _handleShutdown() {
+    if (this.closed) return;
+
+    // Flush synchronously and close
+    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,
+      });
+
+      try {
+        this.socket.write(message + "\n", () => {
+          logger.debug("[reporter] Shutdown flush complete, closing");
+          this.close();
+        });
+      } catch {
+        this.close();
+      }
+    } else {
+      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.5",
   "description": "Token-Optimized Testing Framework for AI-Assisted Development",
   "main": "index.js",
   "type": "module",