pass request id in headers

packages/compass-e2e-tests/helpers/assistant-service.ts

@@ -156,16 +156,16 @@ export async function startMockAssistantServer(
   getResponse: () => MockAssistantResponse;
   setResponse: (response: MockAssistantResponse) => void;
   getRequests: () => {
-    content: any;
-    req: any;
+    content: Record<string, any>;
+    req: http.IncomingMessage;
   }[];
   endpoint: string;
   server: http.Server;
   stop: () => Promise<void>;
 }> {
   let requests: {
-    content: any;
-    req: any;
+    content: Record<string, any>;
+    req: http.IncomingMessage;
   }[] = [];
   let response = _response;
   const server = http

packages/compass-e2e-tests/tests/collection-ai-query.test.ts

@@ -265,6 +265,7 @@ describe('Collection ai query with chatbot (with mocked backend)', function () {
       expect(requests.length).to.equal(1);
 
       const queryRequest = requests[0];
+      expect(queryRequest.req.headers).to.have.property('x-client-request-id');
       // TODO(COMPASS-10125): Switch the model to `mongodb-slim-latest` when
       // enabling this feature.
       expect(queryRequest.content.model).to.equal('mongodb-chat-latest');

packages/compass-generative-ai/src/atlas-ai-service.spec.ts

@@ -1050,8 +1050,13 @@ describe('AtlasAiService', function () {
             expect(fetchStub).to.have.been.calledOnce;
 
             const { args } = fetchStub.firstCall;
-            const requestBody = JSON.parse(args[1].body as string);
 
+            const requestHeaders = args[1].headers as Record<string, string>;
+            expect(requestHeaders['x-client-request-id']).to.equal(
+              input.requestId
+            );
+
+            const requestBody = JSON.parse(args[1].body as string);
             expect(requestBody.model).to.equal('mongodb-chat-latest');
             expect(requestBody.metadata).to.deep.equal({
               userId: '1234',
@@ -1086,6 +1091,7 @@ describe('AtlasAiService', function () {
                   databaseName: 'peanut',
                   requestId: 'abc',
                   signal: new AbortController().signal,
+                  enableStorage: false,
                 },
                 mockConnectionInfo
               );

packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts

@@ -13,6 +13,7 @@ const OPTIONS: PromptContextOptions = {
   collectionName: 'listings',
   userId: 'test-user-id',
   enableStorage: false,
+  requestId: 'test-request-id',
   schema: {
     _id: {
       types: [
@@ -48,6 +49,7 @@ describe('GenAI Prompts', function () {
     );
     expect(metadata.userId).to.equal(OPTIONS.userId);
     expect(metadata.store).to.equal('false');
+    expect(metadata.requestId).to.equal(OPTIONS.requestId);
 
     expect(prompt).to.be.a('string');
     expect(prompt).to.include(
@@ -91,6 +93,7 @@ describe('GenAI Prompts', function () {
     );
     expect(metadata.userId).to.equal(OPTIONS.userId);
     expect(metadata.store).to.equal('false');
+    expect(metadata.requestId).to.equal(OPTIONS.requestId);
 
     expect(prompt).to.be.a('string');
     expect(prompt).to.include(

packages/compass-generative-ai/src/utils/gen-ai-prompt.ts

@@ -69,6 +69,7 @@ type BuildPromptOptions = {
 type BuildMetadataOptions = {
   userId: string;
   enableStorage: boolean;
+  requestId: string;
   type: 'find' | 'aggregate';
 };
 
@@ -167,6 +168,7 @@ export type AiQueryPrompt = {
   metadata: {
     instructions: string;
     userId: string;
+    requestId: string;
   } & (
     | {
         store: 'true';
@@ -181,6 +183,7 @@ export type AiQueryPrompt = {
 function buildMetadata({
   type,
   userId,
+  requestId,
   enableStorage,
 }: BuildMetadataOptions): AiQueryPrompt['metadata'] {
   return {
@@ -189,6 +192,7 @@ function buildMetadata({
         ? buildInstructionsForFindQuery()
         : buildInstructionsForAggregateQuery(),
     userId,
+    requestId,
     ...(enableStorage
       ? {
           sensitiveStorage: 'sensitive',
@@ -203,6 +207,7 @@ function buildMetadata({
 export function buildFindQueryPrompt({
   userId,
   enableStorage,
+  requestId,
   ...restOfTheOptions
 }: PromptContextOptions): AiQueryPrompt {
   const type = 'find';
@@ -215,6 +220,7 @@ export function buildFindQueryPrompt({
     metadata: buildMetadata({
       type,
       userId,
+      requestId,
       enableStorage,
     }),
   };
@@ -223,6 +229,7 @@ export function buildFindQueryPrompt({
 export function buildAggregateQueryPrompt({
   userId,
   enableStorage,
+  requestId,
   ...restOfTheOptions
 }: PromptContextOptions): AiQueryPrompt {
   const type = 'aggregate';
@@ -232,6 +239,6 @@ export function buildAggregateQueryPrompt({
   });
   return {
     prompt,
-    metadata: buildMetadata({ type, userId, enableStorage }),
+    metadata: buildMetadata({ type, userId, requestId, enableStorage }),
   };
 }

packages/compass-generative-ai/src/utils/gen-ai-response.ts

@@ -8,7 +8,7 @@ export async function getAiQueryResponse(
   message: AiQueryPrompt,
   abortSignal: AbortSignal
 ): Promise<string> {
-  const { instructions, ...restOfMetadata } = message.metadata;
+  const { instructions, requestId, ...restOfMetadata } = message.metadata;
   const response = streamText({
     model,
     messages: [{ role: 'user', content: message.prompt }],
@@ -18,6 +18,9 @@ export async function getAiQueryResponse(
         metadata: restOfMetadata,
       },
     },
+    headers: {
+      'X-Client-Request-Id': requestId,
+    },
     abortSignal,
   }).toUIMessageStream();
   const chunks: string[] = [];