taist 0.1.6 → 0.1.8

package/instrument.js

@@ -46,7 +46,8 @@ import {
   instrumentDirectory,
   instrumentModules,
   wrapWithContext,
-  instrumentClassWithContext
+  instrumentClassWithContext,
+  instrumentModule
 } from './lib/instrument-all.js';
 
 // Initialize global reporter for cross-process trace collection
@@ -277,6 +278,7 @@ export {
   instrumentAll,
   instrumentDirectory,
   instrumentModules,
+  instrumentModule,
   wrapWithContext,
   startTrace,
   getContext,

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();
@@ -155,8 +161,10 @@ export class TraceReporter extends EventEmitter {
    * Flushes buffer and waits for data to be sent before closing.
    */
   _handleShutdown() {
-    if (this.closed) return;
+    if (this.closed || this.shuttingDown) return;
 
+    // Set flag immediately to prevent exit handlers from interfering
+    this.shuttingDown = true;
     this._stopFlushTimer();
 
     // Flush and wait for data to be sent

package/package.json

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

package/types/instrument.d.ts

@@ -199,6 +199,28 @@ export declare function instrumentModules(
   options?: InstrumentAllOptions
 ): Promise<Record<string, unknown>>;
 
+/**
+ * Instrument all exports from a module object
+ *
+ * Wraps all function exports with context-aware tracing.
+ * Classes are wrapped so new instances are automatically instrumented.
+ *
+ * @param moduleExports Module exports object (e.g., from `import * as mod from './mod.js'`)
+ * @param moduleName Module name prefix for trace names
+ * @returns Object with same keys but instrumented values
+ *
+ * @example
+ * import { instrumentModule } from 'taist/instrument';
+ * import * as orderServices from './services/order.js';
+ *
+ * export const Order = instrumentModule(orderServices, 'Order');
+ * // All functions in Order will now be traced
+ */
+export declare function instrumentModule<T extends Record<string, unknown>>(
+  moduleExports: T,
+  moduleName: string
+): T;
+
 /**
  * Wrap a function with trace context
  * @param fn Function to wrap