taist 0.2.0 → 0.2.2

package/instrument.js

@@ -163,6 +163,33 @@ export function instrumentExpress(app, options = {}) {
   const useContext = options.useContext !== false;
   const reporter = getGlobalReporter();
 
+  // Add early middleware to set up correlation ID for ALL requests
+  // This runs before any route handlers or other middleware
+  const debug = process.env.TAIST_DEBUG === 'true';
+  app.use((req, res, next) => {
+    // Only set up if not already set (avoid double-instrumentation)
+    if (!req.taistCorrelationId) {
+      const correlationId = generateId();
+      req.taistCorrelationId = correlationId;
+      setCorrelationId(correlationId);
+
+      if (debug) {
+        logger.log('[taist middleware] Set correlationId:', correlationId, 'for', req.method, req.path);
+      }
+
+      // Clear correlation ID when response finishes
+      res.on('finish', () => {
+        if (debug) {
+          logger.log('[taist middleware] Clearing correlationId for', req.method, req.path);
+        }
+        clearCorrelationId();
+      });
+    } else if (debug) {
+      logger.log('[taist middleware] Already has correlationId:', req.taistCorrelationId, 'for', req.method, req.path);
+    }
+    next();
+  });
+
   // Instrument route handlers
   ['get', 'post', 'put', 'delete', 'patch', 'options', 'head'].forEach(method => {
     const original = app[method];
@@ -176,14 +203,15 @@ export function instrumentExpress(app, options = {}) {
         return async function (req, res, next) {
           const traceName = `Route.${method.toUpperCase()} ${path}`;
           const id = generateId();
-          const correlationId = generateId(); // Unique per HTTP request
+          // Use existing correlationId from early middleware, or create one
+          const correlationId = req.taistCorrelationId || generateId();
           const start = performance.now();
 
-          // Store correlationId on req for Apollo/GraphQL integration
-          req.taistCorrelationId = correlationId;
-
-          // Set fallback correlation ID for frameworks that break AsyncLocalStorage
-          setCorrelationId(correlationId);
+          // Ensure correlationId is on req (in case early middleware didn't run)
+          if (!req.taistCorrelationId) {
+            req.taistCorrelationId = correlationId;
+            setCorrelationId(correlationId);
+          }
 
           // Report function for success/error
           const reportResult = (error = null) => {
@@ -217,8 +245,7 @@ export function instrumentExpress(app, options = {}) {
             if (!responseSent) {
               responseSent = true;
               reportResult();
-              // Clear correlation ID after request completes
-              clearCorrelationId();
+              // Note: correlationId is cleared by early middleware's res.on('finish')
             }
             return originalEnd.apply(this, args);
           };
@@ -241,7 +268,6 @@ export function instrumentExpress(app, options = {}) {
               if (!responseSent) {
                 responseSent = true;
                 reportResult(err);
-                clearCorrelationId();
               }
               throw err;
             }
@@ -253,7 +279,6 @@ export function instrumentExpress(app, options = {}) {
               if (!responseSent) {
                 responseSent = true;
                 reportResult(err);
-                clearCorrelationId();
               }
               throw err;
             }
@@ -332,7 +357,16 @@ export function instrumentServiceWithContext(service, name) {
  * };
  */
 export function bridgeContext(req) {
-  const correlationId = req?.taistCorrelationId || getCorrelationId();
+  const debug = process.env.TAIST_DEBUG === 'true';
+  const reqCorrelationId = req?.taistCorrelationId;
+  const fallbackCorrelationId = getCorrelationId();
+  const correlationId = reqCorrelationId || fallbackCorrelationId;
+
+  if (debug) {
+    logger.log('[bridgeContext] req.taistCorrelationId:', reqCorrelationId);
+    logger.log('[bridgeContext] getCorrelationId():', fallbackCorrelationId);
+    logger.log('[bridgeContext] using:', correlationId);
+  }
 
   // Also set the fallback so resolvers can access it even without context prop-drilling
   if (correlationId) {

package/package.json

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