taist 0.2.16 → 0.2.18

package/instrument.js

@@ -92,6 +92,45 @@ export async function flushTraces() {
   }
 }
 
+/**
+ * Gracefully shutdown taist - call this before process exit.
+ * Waits for pending socket writes to complete and flushes remaining traces.
+ *
+ * This is the recommended way to ensure all traces are captured when integrating
+ * with applications that have their own shutdown hooks (e.g., Directus onShutdown).
+ *
+ * @param {number} timeoutMs - Maximum time to wait for pending writes (default: 2000ms)
+ * @returns {Promise<void>}
+ *
+ * @example
+ * // Directus integration
+ * import { gracefulShutdown } from 'taist/instrument';
+ *
+ * export default defineHook(({ onShutdown }) => {
+ *   onShutdown(async () => {
+ *     await gracefulShutdown();
+ *   });
+ * });
+ */
+export async function gracefulShutdown(timeoutMs = 2000) {
+  if (!reporter) return;
+
+  // Prevent SIGTERM handler from doing redundant work
+  if (reporter.shuttingDown || reporter.closed) return;
+  reporter.shuttingDown = true;
+
+  // 1. Wait for pending writes to complete
+  if (reporter._waitForPendingWrites) {
+    await reporter._waitForPendingWrites(timeoutMs);
+  }
+
+  // 2. Final flush of any buffered traces
+  await reporter.flush();
+
+  // 3. Close socket cleanly
+  reporter.close();
+}
+
 // Log initialization
 if (tracer.options.enabled) {
   logger.log('Instrumentation enabled');

package/lib/trace-reporter.js

@@ -37,6 +37,7 @@ export class TraceReporter extends EventEmitter {
     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
+    this._flushScheduled = false; // For micro-batching in flushImmediate mode
 
     logger.debug("[reporter] Created with socketPath:", this.socketPath, "flushImmediate:", this.flushImmediate);
 
@@ -47,9 +48,9 @@ export class TraceReporter extends EventEmitter {
   }
 
   _setupExitHandlers() {
-    const cleanup = (signal) => {
-      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
+    // Synchronous cleanup for exit events
+    const syncCleanup = (signal) => {
+      logger.debug("[reporter] Sync cleanup triggered by:", signal, "buffer size:", this.buffer.length);
       if (this.shuttingDown) {
         logger.debug("[reporter] Shutdown already in progress, skipping cleanup");
         return;
@@ -60,10 +61,75 @@ export class TraceReporter extends EventEmitter {
       }
     };
 
-    process.on("beforeExit", () => cleanup("beforeExit"));
-    process.on("exit", () => cleanup("exit"));
-    process.on("SIGINT", () => cleanup("SIGINT"));
-    process.on("SIGTERM", () => cleanup("SIGTERM"));
+    // Async cleanup for signals - waits for pending writes then exits
+    const signalCleanup = async (signal) => {
+      const lifecycleDebug = process.env.TAIST_TRACE_LIFECYCLE === 'true';
+      if (lifecycleDebug) {
+        console.log(`[LIFECYCLE reporter] Signal ${signal} received, pendingWrites:`, this.pendingWrites);
+      }
+
+      if (this.shuttingDown) {
+        logger.debug("[reporter] Shutdown already in progress, skipping cleanup");
+        return;
+      }
+      this.shuttingDown = true;
+
+      // Wait for pending writes to complete (up to 2 seconds)
+      await this._waitForPendingWrites(2000);
+
+      if (lifecycleDebug) {
+        console.log(`[LIFECYCLE reporter] After waiting, pendingWrites:`, this.pendingWrites, 'buffer:', this.buffer.length);
+      }
+
+      // Final flush and close
+      if (!this.closed) {
+        this.flushSync();
+        this.close();
+      }
+
+      // Exit the process after cleanup
+      process.exit(0);
+    };
+
+    process.on("beforeExit", () => syncCleanup("beforeExit"));
+    process.on("exit", () => syncCleanup("exit"));
+    process.on("SIGINT", () => signalCleanup("SIGINT"));
+    process.on("SIGTERM", () => signalCleanup("SIGTERM"));
+  }
+
+  /**
+   * Wait for pending socket writes to complete.
+   * @param {number} timeout - Maximum time to wait in milliseconds
+   */
+  async _waitForPendingWrites(timeout = 2000) {
+    if (this.pendingWrites === 0) return;
+
+    const lifecycleDebug = process.env.TAIST_TRACE_LIFECYCLE === 'true';
+    const startTime = Date.now();
+
+    return new Promise((resolve) => {
+      const checkPending = () => {
+        if (this.pendingWrites === 0) {
+          if (lifecycleDebug) {
+            console.log('[LIFECYCLE reporter] All pending writes completed');
+          }
+          resolve();
+          return;
+        }
+
+        if ((Date.now() - startTime) >= timeout) {
+          if (lifecycleDebug) {
+            console.log('[LIFECYCLE reporter] Timeout waiting for pending writes, still have:', this.pendingWrites);
+          }
+          resolve();
+          return;
+        }
+
+        // Poll at 10ms intervals
+        setTimeout(checkPending, 10);
+      };
+      checkPending();
+    });
   }
 
   async connect() {
@@ -313,10 +379,18 @@ export class TraceReporter extends EventEmitter {
       });
     }
 
-    // Flush immediately if configured (for spawned processes with unpredictable exit)
-    // Otherwise auto-flush when batch size reached
+    // Flush strategy:
+    // - flushImmediate=true: Use micro-batching (setImmediate) to batch traces from the same tick
+    //   This significantly reduces socket writes (e.g., 2762 individual writes → ~100 batched writes)
+    // - flushImmediate=false: Traditional batching by size (flushes when buffer reaches batchSize)
     if (this.flushImmediate) {
-      this.flush().catch(() => {});
+      if (!this._flushScheduled) {
+        this._flushScheduled = true;
+        setImmediate(() => {
+          this._flushScheduled = false;
+          this.flush().catch(() => {});
+        });
+      }
     } else if (this.buffer.length >= this.batchSize) {
       this.flush().catch(() => {});
     }

package/package.json

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

package/types/instrument.d.ts

@@ -111,6 +111,28 @@ export declare const reporter: TraceReporter;
  */
 export declare function flushTraces(): Promise<void>;
 
+/**
+ * Gracefully shutdown taist - call this before process exit.
+ * Waits for pending socket writes to complete and flushes remaining traces.
+ *
+ * This is the recommended way to ensure all traces are captured when integrating
+ * with applications that have their own shutdown hooks (e.g., Directus onShutdown).
+ *
+ * @param timeoutMs Maximum time to wait for pending writes (default: 2000ms)
+ * @returns Promise that resolves when shutdown is complete
+ *
+ * @example
+ * // Directus integration
+ * import { gracefulShutdown } from 'taist/instrument';
+ *
+ * export default defineHook(({ onShutdown }) => {
+ *   onShutdown(async () => {
+ *     await gracefulShutdown();
+ *   });
+ * });
+ */
+export declare function gracefulShutdown(timeoutMs?: number): Promise<void>;
+
 /**
  * Auto-instrument a module's exports
  * @param moduleExports Module exports object