Capture console logs in Koa middleware

Author: itssimon

src/koa/middleware.ts

@@ -1,9 +1,15 @@
 import Koa from "koa";
 
+import { AsyncLocalStorage } from "async_hooks";
 import { ApitallyClient } from "../common/client.js";
+import { patchConsole } from "../common/consoleCapture.js";
 import { consumerFromStringOrObject } from "../common/consumerRegistry.js";
 import { getPackageVersion } from "../common/packageVersions.js";
-import { convertBody, convertHeaders } from "../common/requestLogger.js";
+import {
+  convertBody,
+  convertHeaders,
+  LogRecord,
+} from "../common/requestLogger.js";
 import {
   ApitallyConfig,
   ApitallyConsumer,
@@ -28,87 +34,97 @@ export function useApitally(app: Koa, config: ApitallyConfig) {
 }
 
 function getMiddleware(client: ApitallyClient) {
+  const logsContext = new AsyncLocalStorage<LogRecord[]>();
+
+  if (client.requestLogger.config.captureLogs) {
+    patchConsole(logsContext);
+  }
+
   return async (ctx: Koa.Context, next: Koa.Next) => {
     if (!client.isEnabled() || ctx.request.method.toUpperCase() === "OPTIONS") {
       await next();
       return;
     }
 
-    let path: string | undefined;
-    let statusCode: number | undefined;
-    let serverError: Error | undefined;
-    const startTime = performance.now();
-    try {
-      await next();
-    } catch (error: any) {
-      path = getPath(ctx);
-      statusCode = error.statusCode || error.status || 500;
-      if (path && statusCode === 500 && error instanceof Error) {
-        serverError = error;
-        client.serverErrorCounter.addServerError({
-          consumer: getConsumer(ctx)?.identifier,
-          method: ctx.request.method,
-          path,
-          type: error.name,
-          msg: error.message,
-          traceback: error.stack || "",
-        });
-      }
-      throw error;
-    } finally {
-      const responseTime = performance.now() - startTime;
-      const consumer = getConsumer(ctx);
-      client.consumerRegistry.addOrUpdateConsumer(consumer);
-      if (!path) {
+    await logsContext.run([], async () => {
+      let path: string | undefined;
+      let statusCode: number | undefined;
+      let serverError: Error | undefined;
+      const startTime = performance.now();
+      try {
+        await next();
+      } catch (error: any) {
         path = getPath(ctx);
-      }
-      if (path) {
-        try {
-          client.requestCounter.addRequest({
-            consumer: consumer?.identifier,
+        statusCode = error.statusCode || error.status || 500;
+        if (path && statusCode === 500 && error instanceof Error) {
+          serverError = error;
+          client.serverErrorCounter.addServerError({
+            consumer: getConsumer(ctx)?.identifier,
             method: ctx.request.method,
             path,
-            statusCode: statusCode || ctx.response.status,
-            responseTime,
-            requestSize: ctx.request.length,
-            responseSize: ctx.response.length,
+            type: error.name,
+            msg: error.message,
+            traceback: error.stack || "",
           });
-        } catch (error) {
-          client.logger.error(
-            "Error while logging request in Apitally middleware.",
-            { context: ctx, error },
+        }
+        throw error;
+      } finally {
+        const responseTime = performance.now() - startTime;
+        const consumer = getConsumer(ctx);
+        client.consumerRegistry.addOrUpdateConsumer(consumer);
+        if (!path) {
+          path = getPath(ctx);
+        }
+        if (path) {
+          try {
+            client.requestCounter.addRequest({
+              consumer: consumer?.identifier,
+              method: ctx.request.method,
+              path,
+              statusCode: statusCode || ctx.response.status,
+              responseTime,
+              requestSize: ctx.request.length,
+              responseSize: ctx.response.length,
+            });
+          } catch (error) {
+            client.logger.error(
+              "Error while logging request in Apitally middleware.",
+              { context: ctx, error },
+            );
+          }
+        }
+        if (client.requestLogger.enabled) {
+          const logs = logsContext.getStore();
+          client.requestLogger.logRequest(
+            {
+              timestamp: Date.now() / 1000,
+              method: ctx.request.method,
+              path,
+              url: ctx.request.href,
+              headers: convertHeaders(ctx.request.headers),
+              size: ctx.request.length,
+              consumer: consumer?.identifier,
+              body: convertBody(
+                ctx.request.body,
+                ctx.request.get("content-type"),
+              ),
+            },
+            {
+              statusCode: statusCode || ctx.response.status,
+              responseTime: responseTime / 1000,
+              headers: convertHeaders(ctx.response.headers),
+              size: ctx.response.length,
+              body: convertBody(
+                ctx.response.body,
+                ctx.response.get("content-type"),
+              ),
+            },
+            serverError,
+            logs,
           );
         }
       }
-      if (client.requestLogger.enabled) {
-        client.requestLogger.logRequest(
-          {
-            timestamp: Date.now() / 1000,
-            method: ctx.request.method,
-            path,
-            url: ctx.request.href,
-            headers: convertHeaders(ctx.request.headers),
-            size: ctx.request.length,
-            consumer: consumer?.identifier,
-            body: convertBody(
-              ctx.request.body,
-              ctx.request.get("content-type"),
-            ),
-          },
-          {
-            statusCode: statusCode || ctx.response.status,
-            responseTime: responseTime / 1000,
-            headers: convertHeaders(ctx.response.headers),
-            size: ctx.response.length,
-            body: convertBody(
-              ctx.response.body,
-              ctx.response.get("content-type"),
-            ),
-          },
-          serverError,
-        );
-      }
-    }
+    });
   };
 }
 

tests/koa/app.test.ts

@@ -97,6 +97,12 @@ testCases.forEach(({ name, router, getApp }) => {
       ]);
       expect(call[1].body).toBeInstanceOf(Buffer);
       expect(call[1].body!.toString()).toMatch(/^Hello John!/);
+      expect(call[3]).toBeDefined();
+      expect(call[3]).toHaveLength(2);
+      expect(call[3]![0].level).toBe("log");
+      expect(call[3]![0].message).toBe("Test 1");
+      expect(call[3]![1].level).toBe("warn");
+      expect(call[3]![1].message).toBe("Test 2");
       spy.mockReset();
 
       await appTest.post("/hello").send({ name: "John", age: 20 }).expect(200);

tests/koa/app.ts

@@ -13,6 +13,7 @@ const requestLoggingConfig = {
   logRequestBody: true,
   logResponseHeaders: true,
   logResponseBody: true,
+  captureLogs: true,
 };
 
 export const getAppWithKoaRouter = () => {
@@ -28,6 +29,8 @@ export const getAppWithKoaRouter = () => {
 
   router.get("/hello", async (ctx) => {
     setConsumer(ctx, "test");
+    console.log("Test 1");
+    console.warn("Test 2");
     ctx.body = `Hello ${ctx.query.name}! You are ${ctx.query.age} years old!`;
   });
   router.get("/hello/:id", async (ctx) => {
@@ -62,6 +65,8 @@ export const getAppWithKoaRoute = () => {
   app.use(
     route.get("/hello", async (ctx) => {
       setConsumer(ctx, "test");
+      console.log("Test 1");
+      console.warn("Test 2");
       ctx.body = `Hello ${ctx.query.name}! You are ${ctx.query.age} years old!`;
     }),
   );