taist 0.1.5 → 0.1.7

package/lib/trace-collector.js

@@ -184,28 +184,44 @@ export class TraceCollector extends EventEmitter {
     }
 
     return new Promise((resolve) => {
-      // Force close any remaining 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

@@ -30,6 +30,7 @@ export class TraceReporter extends EventEmitter {
     this.flushTimer = null;
     this.workerId = options.workerId || process.pid;
     this.closed = false;
+    this.shuttingDown = false; // Prevents race between shutdown signal and SIGTERM
 
     logger.debug("[reporter] Created with socketPath:", this.socketPath);
 
@@ -41,7 +42,12 @@ export class TraceReporter extends EventEmitter {
 
   _setupExitHandlers() {
     const cleanup = (signal) => {
-      logger.debug("[reporter] Cleanup triggered by:", signal, "buffer size:", this.buffer.length);
+      logger.debug("[reporter] Cleanup triggered by:", signal, "buffer size:", this.buffer.length, "shuttingDown:", this.shuttingDown);
+      // If shutdown signal handler is already handling graceful shutdown, don't interfere
+      if (this.shuttingDown) {
+        logger.debug("[reporter] Shutdown already in progress, skipping cleanup");
+        return;
+      }
       if (!this.closed) {
         this.flushSync();
         this.close();
@@ -152,12 +158,16 @@ export class TraceReporter extends EventEmitter {
 
   /**
    * Handle shutdown signal from collector.
-   * Immediately flushes buffer and closes connection.
+   * Flushes buffer and waits for data to be sent before closing.
    */
   _handleShutdown() {
-    if (this.closed) return;
+    if (this.closed || this.shuttingDown) return;
 
-    // Flush synchronously and close
+    // 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({
@@ -166,15 +176,63 @@ export class TraceReporter extends EventEmitter {
         data: traces,
       });
 
+      logger.debug("[reporter] Shutdown flushing", traces.length, "traces");
+
       try {
-        this.socket.write(message + "\n", () => {
-          logger.debug("[reporter] Shutdown flush complete, closing");
-          this.close();
-        });
-      } catch {
+        // 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();
     }
   }

package/package.json

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