chore: abort cursor iteration on request timeouts

himanshusinghs · himanshusinghs · commit 13d840812833 · 2025-09-10T20:32:43.000+02:00
diff --git a/src/helpers/iterateCursor.ts b/src/helpers/iterateCursor.ts
@@ -8,10 +8,15 @@ import type { AggregationCursor, FindCursor } from "mongodb";
  * The cursor is iterated until we can predict that fetching next doc won't
  * exceed the maxBytesPerQuery limit.
  */
-export async function iterateCursorUntilMaxBytes(
-    cursor: FindCursor<unknown> | AggregationCursor<unknown>,
-    maxBytesPerQuery: number
-): Promise<unknown[]> {
+export async function iterateCursorUntilMaxBytes({
+    cursor,
+    maxBytesPerQuery,
+    abortSignal,
+}: {
+    cursor: FindCursor<unknown> | AggregationCursor<unknown>;
+    maxBytesPerQuery: number;
+    abortSignal?: AbortSignal;
+}): Promise<unknown[]> {
     // Setting configured limit to zero or negative is equivalent to disabling
     // the max bytes limit applied on tool responses.
     if (maxBytesPerQuery <= 0) {
@@ -22,6 +27,10 @@ export async function iterateCursorUntilMaxBytes(
     let totalBytes = 0;
     const bufferedDocuments: unknown[] = [];
     while (true) {
+        if (abortSignal?.aborted) {
+            break;
+        }
+
         if (totalBytes + biggestDocSizeSoFar >= maxBytesPerQuery) {
             break;
         }
diff --git a/src/tools/mongodb/read/aggregate.ts b/src/tools/mongodb/read/aggregate.ts
@@ -3,7 +3,7 @@ import type { AggregationCursor } from "mongodb";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
+import type { ToolArgs, OperationType, ToolExecutionContext } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
 import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 import { type Document, EJSON } from "bson";
@@ -25,11 +25,10 @@ export class AggregateTool extends MongoDBToolBase {
     };
     public operationType: OperationType = "read";
 
-    protected async execute({
-        database,
-        collection,
-        pipeline,
-    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+    protected async execute(
+        { database, collection, pipeline }: ToolArgs<typeof this.argsShape>,
+        { signal }: ToolExecutionContext
+    ): Promise<CallToolResult> {
         let aggregationCursor: AggregationCursor | undefined;
         try {
             const provider = await this.ensureConnected();
@@ -53,7 +52,11 @@ export class AggregateTool extends MongoDBToolBase {
 
             const [totalDocuments, documents] = await Promise.all([
                 this.countAggregationResultDocuments({ provider, database, collection, pipeline }),
-                iterateCursorUntilMaxBytes(aggregationCursor, this.config.maxBytesPerQuery),
+                iterateCursorUntilMaxBytes({
+                    cursor: aggregationCursor,
+                    maxBytesPerQuery: this.config.maxDocumentsPerQuery,
+                    abortSignal: signal,
+                }),
             ]);
 
             let messageDescription = `\
diff --git a/src/tools/mongodb/read/find.ts b/src/tools/mongodb/read/find.ts
@@ -1,7 +1,7 @@
 import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
+import type { ToolArgs, OperationType, ToolExecutionContext } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
 import type { FindCursor, SortDirection } from "mongodb";
 import { checkIndexUsage } from "../../../helpers/indexCheck.js";
@@ -40,14 +40,10 @@ export class FindTool extends MongoDBToolBase {
     };
     public operationType: OperationType = "read";
 
-    protected async execute({
-        database,
-        collection,
-        filter,
-        projection,
-        limit,
-        sort,
-    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+    protected async execute(
+        { database, collection, filter, projection, limit, sort }: ToolArgs<typeof this.argsShape>,
+        { signal }: ToolExecutionContext
+    ): Promise<CallToolResult> {
         let findCursor: FindCursor<unknown> | undefined;
         try {
             const provider = await this.ensureConnected();
@@ -82,7 +78,11 @@ export class FindTool extends MongoDBToolBase {
                         }),
                     undefined
                 ),
-                iterateCursorUntilMaxBytes(findCursor, this.config.maxBytesPerQuery),
+                iterateCursorUntilMaxBytes({
+                    cursor: findCursor,
+                    maxBytesPerQuery: this.config.maxBytesPerQuery,
+                    abortSignal: signal,
+                }),
             ]);
 
             let messageDescription = `\
diff --git a/src/tools/tool.ts b/src/tools/tool.ts
@@ -11,6 +11,8 @@ import type { Server } from "../server.js";
 
 export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNever>;
 
+export type ToolExecutionContext<Args extends ZodRawShape = ZodRawShape> = Parameters<ToolCallback<Args>>[1];
+
 export type OperationType = "metadata" | "read" | "create" | "delete" | "update" | "connect";
 export type ToolCategory = "mongodb" | "atlas";
 export type TelemetryToolMetadata = {
diff --git a/tests/unit/helpers/iterateCursor.test.ts b/tests/unit/helpers/iterateCursor.test.ts
@@ -4,10 +4,17 @@ import { calculateObjectSize } from "bson";
 import { iterateCursorUntilMaxBytes } from "../../../src/helpers/iterateCursor.js";
 
 describe("iterateCursorUntilMaxBytes", () => {
-    function createMockCursor(docs: unknown[]): FindCursor<unknown> {
+    function createMockCursor(
+        docs: unknown[],
+        { abortController, abortOnIdx }: { abortController?: AbortController; abortOnIdx?: number } = {}
+    ): FindCursor<unknown> {
         let idx = 0;
         return {
             tryNext: vi.fn(() => {
+                if (idx === abortOnIdx) {
+                    abortController?.abort();
+                }
+
                 if (idx < docs.length) {
                     return Promise.resolve(docs[idx++]);
                 }
@@ -23,23 +30,36 @@ describe("iterateCursorUntilMaxBytes", () => {
         const docs = Array.from({ length: 1000 }).map((_, idx) => ({ value: idx }));
         const cursor = createMockCursor(docs);
         const maxBytes = -1;
-        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        const result = await iterateCursorUntilMaxBytes({ cursor, maxBytesPerQuery: maxBytes });
         expect(result).toEqual(docs);
     });
 
     it("returns all docs if maxBytesPerQuery is 0", async () => {
         const docs = Array.from({ length: 1000 }).map((_, idx) => ({ value: idx }));
         const cursor = createMockCursor(docs);
         const maxBytes = 0;
-        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        const result = await iterateCursorUntilMaxBytes({ cursor, maxBytesPerQuery: maxBytes });
         expect(result).toEqual(docs);
     });
 
+    it("respects abort signal and breaks out of loop when aborted", async () => {
+        const docs = Array.from({ length: 20 }).map((_, idx) => ({ value: idx }));
+        const abortController = new AbortController();
+        const cursor = createMockCursor(docs, { abortOnIdx: 9, abortController });
+        const maxBytes = 10000;
+        const result = await iterateCursorUntilMaxBytes({
+            cursor,
+            maxBytesPerQuery: maxBytes,
+            abortSignal: abortController.signal,
+        });
+        expect(result).toEqual(Array.from({ length: 10 }).map((_, idx) => ({ value: idx })));
+    });
+
     it("returns all docs if under maxBytesPerQuery", async () => {
         const docs = [{ a: 1 }, { b: 2 }];
         const cursor = createMockCursor(docs);
         const maxBytes = 10000;
-        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        const result = await iterateCursorUntilMaxBytes({ cursor, maxBytesPerQuery: maxBytes });
         expect(result).toEqual(docs);
     });
 
@@ -49,20 +69,20 @@ describe("iterateCursorUntilMaxBytes", () => {
         const docs = [doc1, doc2];
         const cursor = createMockCursor(docs);
         const maxBytes = calculateObjectSize(doc1) + 10;
-        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        const result = await iterateCursorUntilMaxBytes({ cursor, maxBytesPerQuery: maxBytes });
         expect(result).toEqual([doc1]);
     });
 
     it("returns empty array if maxBytesPerQuery is smaller than even the first doc", async () => {
         const docs = [{ a: "x".repeat(100) }];
         const cursor = createMockCursor(docs);
-        const result = await iterateCursorUntilMaxBytes(cursor, 10);
+        const result = await iterateCursorUntilMaxBytes({ cursor, maxBytesPerQuery: 10 });
         expect(result).toEqual([]);
     });
 
     it("handles empty cursor", async () => {
         const cursor = createMockCursor([]);
-        const result = await iterateCursorUntilMaxBytes(cursor, 1000);
+        const result = await iterateCursorUntilMaxBytes({ cursor, maxBytesPerQuery: 1000 });
         expect(result).toEqual([]);
     });
 
@@ -73,7 +93,7 @@ describe("iterateCursorUntilMaxBytes", () => {
         const cursor = createMockCursor(docs);
         // Set maxBytes so that after doc1, biggestDocSizeSoFar would prevent fetching doc2
         const maxBytes = calculateObjectSize(doc1) + calculateObjectSize(doc2) - 1;
-        const result = await iterateCursorUntilMaxBytes(cursor, maxBytes);
+        const result = await iterateCursorUntilMaxBytes({ cursor, maxBytesPerQuery: maxBytes });
         // Should only include doc1, not doc2
         expect(result).toEqual([doc1]);
     });
