(EAI-1159) Implement middleware in new contentRouter

packages/chatbot-server-mongodb-public/src/config.ts

@@ -16,9 +16,10 @@ import {
   requireRequestOrigin,
   AddCustomDataFunc,
   makeDefaultFindVerifiedAnswer,
-  defaultCreateConversationCustomData,
-  defaultAddMessageToConversationCustomData,
   makeVerifiedAnswerGenerateResponse,
+  addDefaultCustomData,
+  ConversationsRouterLocals,
+  SearchContentRouterLocals,
 } from "mongodb-chatbot-server";
 import cookieParser from "cookie-parser";
 import { blockGetRequests } from "./middleware/blockGetRequests";
@@ -269,7 +270,7 @@ export const generateResponse = wrapTraced(
 
 export const createConversationCustomDataWithAuthUser: AddCustomDataFunc =
   async (req, res) => {
-    const customData = await defaultCreateConversationCustomData(req, res);
+    const customData = await addDefaultCustomData(req, res);
     if (req.cookies.auth_user) {
       customData.authUser = req.cookies.auth_user;
     }
@@ -318,12 +319,13 @@ export const config: AppConfig = {
       classifierModel: languageModel,
     }),
     searchResultsStore,
+    middleware: [requireValidIpAddress<SearchContentRouterLocals>(), requireRequestOrigin<SearchContentRouterLocals>()],
   },
   conversationsRouterConfig: {
     middleware: [
       blockGetRequests,
-      requireValidIpAddress(),
-      requireRequestOrigin(),
+      requireValidIpAddress<ConversationsRouterLocals>(),
+      requireRequestOrigin<ConversationsRouterLocals>(),
       useSegmentIds(),
       redactConnectionUri(),
       cookieParser(),
@@ -332,10 +334,7 @@ export const config: AppConfig = {
       ? createConversationCustomDataWithAuthUser
       : undefined,
     addMessageToConversationCustomData: async (req, res) => {
-      const defaultCustomData = await defaultAddMessageToConversationCustomData(
-        req,
-        res
-      );
+      const defaultCustomData = await addDefaultCustomData(req, res);
       const customData = {
         ...defaultCustomData,
       };

packages/chatbot-server-mongodb-public/src/processors/findContentWithMongoDbMetadata.ts

@@ -31,7 +31,7 @@ export const makeFindContentWithMongoDbMetadata = ({
       return res;
     },
     {
-      name: "makeFindContentWithMongoDbMetadata",
+      name: "findContentWithMongoDbMetadata",
     }
   );
   return wrappedFindContent;

packages/mongodb-chatbot-server/src/middleware/requireRequestOrigin.ts

@@ -1,9 +1,13 @@
+import { RequestHandler } from "express";
+import { ParamsDictionary } from "express-serve-static-core";
+import { ParsedQs } from "qs";
 import { getRequestId, logRequest, sendErrorResponse } from "../utils";
-import { ConversationsMiddleware } from "../routes/conversations/conversationsRouter";
 
 export const CUSTOM_REQUEST_ORIGIN_HEADER = "X-Request-Origin";
 
-export function requireRequestOrigin(): ConversationsMiddleware {
+export function requireRequestOrigin<
+  Locals extends Record<string, any>
+>(): RequestHandler<ParamsDictionary, unknown, unknown, ParsedQs, Locals> {
   return (req, res, next) => {
     const reqId = getRequestId(req);
 

packages/mongodb-chatbot-server/src/middleware/requireValidIpAddress.test.ts

@@ -3,6 +3,7 @@ import {
   createConversationsMiddlewareReq,
   createConversationsMiddlewareRes,
 } from "../test/middlewareTestHelpers";
+import { ConversationsRouterLocals } from "../routes";
 
 const baseReq = {
   body: { message: "Hello, world!" },
@@ -18,7 +19,7 @@ describe("requireValidIpAddress", () => {
     const res = createConversationsMiddlewareRes();
     const next = jest.fn();
 
-    const middleware = requireValidIpAddress();
+    const middleware = requireValidIpAddress<ConversationsRouterLocals>();
     req.body = baseReq.body;
     req.params = baseReq.params;
     req.query = baseReq.query;
@@ -39,7 +40,7 @@ describe("requireValidIpAddress", () => {
     const next = jest.fn();
 
     const invalidIpAddress = "not-an-ip-address";
-    const middleware = requireValidIpAddress();
+    const middleware = requireValidIpAddress<ConversationsRouterLocals>();
     req.body = baseReq.body;
     req.params = baseReq.params;
     req.query = baseReq.query;
@@ -59,7 +60,7 @@ describe("requireValidIpAddress", () => {
     const res = createConversationsMiddlewareRes();
     const next = jest.fn();
 
-    const middleware = requireValidIpAddress();
+    const middleware = requireValidIpAddress<ConversationsRouterLocals>();
     req.body = baseReq.body;
     req.params = baseReq.params;
     req.query = baseReq.query;

packages/mongodb-chatbot-server/src/middleware/requireValidIpAddress.ts

@@ -1,8 +1,12 @@
+import { RequestHandler } from "express";
+import { ParamsDictionary } from "express-serve-static-core";
+import { ParsedQs } from "qs";
 import { getRequestId, logRequest, sendErrorResponse } from "../utils";
-import { ConversationsMiddleware } from "../routes/conversations/conversationsRouter";
 import { isValidIp } from "../routes/conversations/utils";
 
-export function requireValidIpAddress(): ConversationsMiddleware {
+export function requireValidIpAddress<
+  Locals extends Record<string, any>
+>(): RequestHandler<ParamsDictionary, unknown, unknown, ParsedQs, Locals> {
   return (req, res, next) => {
     const reqId = getRequestId(req);
 

packages/mongodb-chatbot-server/src/middleware/validateRequestSchema.ts

@@ -5,12 +5,12 @@ import { getRequestId, logRequest, sendErrorResponse } from "../utils";
 
 export const SomeExpressRequest = z.object({
   headers: z.object({}).optional(),
-  params: z.object({}).optional(),
+  params: z.object({}),
   query: z.object({}).optional(),
   body: z.object({}).optional(),
 });
 
-function generateZodErrorMessage(error: ZodError) {
+export function generateZodErrorMessage(error: ZodError) {
   return generateErrorMessage(error.issues, {
     delimiter: {
       error: "\n",

packages/mongodb-chatbot-server/src/processors/addCustomData.ts

@@ -0,0 +1,94 @@
+import { Request, Response } from "express";
+
+export type RequestCustomData = Record<string, unknown> | undefined;
+
+/**
+  Function to add custom data to the {@link Conversation} or content search Request persisted to the database.
+  Has access to the Express.js request and response plus the values
+  from the {@link Response.locals} object.
+ */
+export type AddCustomDataFunc = (
+  request: Request,
+  response: Response
+) => Promise<RequestCustomData>;
+
+const addIpToCustomData: AddCustomDataFunc = async (req) =>
+  req.ip
+    ? {
+        ip: req.ip,
+      }
+    : undefined;
+
+const addOriginToCustomData: AddCustomDataFunc = async (_, res) =>
+  res.locals.customData.origin
+    ? {
+        origin: res.locals.customData.origin,
+      }
+    : undefined;
+
+export const originCodes = [
+  "LEARN",
+  "DEVELOPER",
+  "DOCS",
+  "DOTCOM",
+  "GEMINI_CODE_ASSIST",
+  "VSCODE",
+  "OTHER",
+] as const;
+
+export type OriginCode = (typeof originCodes)[number];
+
+interface OriginRule {
+  regex: RegExp;
+  code: OriginCode;
+}
+
+const ORIGIN_RULES: OriginRule[] = [
+  { regex: /learn\.mongodb\.com/, code: "LEARN" },
+  { regex: /mongodb\.com\/developer/, code: "DEVELOPER" },
+  { regex: /mongodb\.com\/docs/, code: "DOCS" },
+  { regex: /mongodb\.com\//, code: "DOTCOM" },
+  { regex: /google-gemini-code-assist/, code: "GEMINI_CODE_ASSIST" },
+  { regex: /vscode-mongodb-copilot/, code: "VSCODE" },
+];
+
+function getOriginCode(origin: string): OriginCode {
+  for (const rule of ORIGIN_RULES) {
+    if (rule.regex.test(origin)) {
+      return rule.code;
+    }
+  }
+  return "OTHER";
+}
+
+const addOriginCodeToCustomData: AddCustomDataFunc = async (_, res) => {
+  const origin = res.locals.customData.origin;
+  return typeof origin === "string" && origin.length > 0
+    ? {
+        originCode: getOriginCode(origin),
+      }
+    : undefined;
+};
+
+const addUserAgentToCustomData: AddCustomDataFunc = async (req) =>
+  req.headers["user-agent"]
+    ? {
+        userAgent: req.headers["user-agent"],
+      }
+    : undefined;
+
+export type AddDefinedCustomDataFunc = (
+  ...args: Parameters<AddCustomDataFunc>
+) => Promise<Exclude<RequestCustomData, undefined>>;
+
+export const addDefaultCustomData: AddDefinedCustomDataFunc = async (
+  req,
+  res
+) => {
+  return {
+    ...(await addIpToCustomData(req, res)),
+    ...(await addOriginToCustomData(req, res)),
+    ...(await addOriginCodeToCustomData(req, res)),
+    ...(await addUserAgentToCustomData(req, res)),
+  };
+};

packages/mongodb-chatbot-server/src/processors/index.ts

@@ -9,3 +9,4 @@ export * from "./InputGuardrail";
 export * from "./makeVerifiedAnswerGenerateResponse";
 export * from "./includeChunksForMaxTokensPossible";
 export * from "./GenerateResponse";
+export * from "./addCustomData";

packages/mongodb-chatbot-server/src/routes/content/contentRouter.test.ts

@@ -1,10 +1,14 @@
+import { Express } from "express";
 import request from "supertest";
-import { makeTestApp } from "../../test/testHelpers";
-import type { MakeContentRouterParams } from "./contentRouter";
 import type {
   FindContentFunc,
   MongoDbSearchResultsStore,
 } from "mongodb-rag-core";
+import type {
+  MakeContentRouterParams,
+  SearchContentMiddleware,
+} from "./contentRouter";
+import { makeTestApp } from "../../test/testHelpers";
 
 // Minimal in-memory mock for SearchResultsStore for testing purposes
 const mockSearchResultsStore: MongoDbSearchResultsStore = {
@@ -23,8 +27,7 @@ const findContentMock = jest.fn().mockResolvedValue({
   queryEmbedding: [],
 }) satisfies FindContentFunc;
 
-// Helper to build contentRouterConfig for the test app
-function makeContentRouterConfig(
+function makeMockContentRouterConfig(
   overrides: Partial<MakeContentRouterParams> = {}
 ) {
   return {
@@ -35,20 +38,57 @@ function makeContentRouterConfig(
 }
 
 describe("contentRouter", () => {
+  const ipAddress = "127.0.0.1";
   const searchEndpoint = "/api/v1/content/search";
 
   it("should call custom middleware if provided", async () => {
     const mockMiddleware = jest.fn((_req, _res, next) => next());
     const { app, origin } = await makeTestApp({
-      contentRouterConfig: makeContentRouterConfig({
+      contentRouterConfig: makeMockContentRouterConfig({
         middleware: [mockMiddleware],
       }),
     });
-    await request(app)
-      .post(searchEndpoint)
-      .set("req-id", "test-req-id")
-      .set("Origin", origin)
-      .send({ query: "mongodb" });
+    await createContentReq({ app, origin, query: "mongodb" });
     expect(mockMiddleware).toHaveBeenCalled();
   });
+
+  test("should use route middleware customData", async () => {
+    const middleware1: SearchContentMiddleware = (_, res, next) => {
+      res.locals.customData.middleware1 = true;
+      next();
+    };
+    let called = false;
+    const middleware2: SearchContentMiddleware = (_, res, next) => {
+      expect(res.locals.customData.middleware1).toBe(true);
+      called = true;
+      next();
+    };
+    const { app, origin } = await makeTestApp({
+      contentRouterConfig: makeMockContentRouterConfig({
+        middleware: [middleware1, middleware2],
+      }),
+    });
+    await createContentReq({ app, origin, query: "What is aggregation?" });
+    expect(called).toBe(true);
+  });
+
+  /**
+    Helper function to create a new content request
+   */
+  async function createContentReq({
+    app,
+    origin,
+    query,
+  }: {
+    app: Express;
+    origin: string;
+    query: string;
+  }) {
+    const createContentRes = await request(app)
+      .post(searchEndpoint)
+      .set("X-FORWARDED-FOR", ipAddress)
+      .set("Origin", origin)
+      .send({ query });
+    return createContentRes;
+  }
 });