Skip to content

Commit da5b4e0

Browse files
nirinchevnirinchev
committed
feat: add structured content to insertmany
1 parent e12068a commit da5b4e0

File tree

4 files changed

+66
-22
lines changed

4 files changed

+66
-22
lines changed

‎src/tools/mongodb/create/insertMany.ts‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { z } from "zod";
2-
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
32
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
3+
import type { ToolResult } from "../../tool.js";
44
import { type ToolArgs, type OperationType, formatUntrustedData } from "../../tool.js";
55
import { zEJSON } from "../../args.js";
66
import { type Document } from "bson";
@@ -37,14 +37,21 @@ export class InsertManyTool extends MongoDBToolBase {
3737
),
3838
}
3939
: commonArgs;
40+
41+
protected outputShape = {
42+
success: z.boolean(),
43+
insertedCount: z.number(),
44+
insertedIds: z.array(z.any()),
45+
};
46+
4047
static operationType: OperationType = "create";
4148

4249
protected async execute({
4350
database,
4451
collection,
4552
documents,
4653
embeddingParameters: providedEmbeddingParameters,
47-
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
54+
}: ToolArgs<typeof this.argsShape>): Promise<ToolResult<typeof this.outputShape>> {
4855
const provider = await this.ensureConnected();
4956

5057
const embeddingParameters = this.isFeatureEnabled("search")
@@ -70,8 +77,14 @@ export class InsertManyTool extends MongoDBToolBase {
7077
`Inserted \`${result.insertedCount}\` document(s) into ${database}.${collection}.`,
7178
`Inserted IDs: ${Object.values(result.insertedIds).join(", ")}`
7279
);
80+
7381
return {
7482
content,
83+
structuredContent: {
84+
success: true,
85+
insertedCount: result.insertedCount,
86+
insertedIds: Object.values(result.insertedIds),
87+
},
7588
};
7689
}
7790

‎src/tools/tool.ts‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import type { z } from "zod";
2-
import { type ZodRawShape, type ZodNever } from "zod";
1+
import type { z, ZodTypeAny } from "zod";
2+
import { type ZodRawShape } from "zod";
33
import type { RegisteredTool, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
44
import type { CallToolResult, ToolAnnotations } from "@modelcontextprotocol/sdk/types.js";
55
import type { Session } from "../common/session.js";
@@ -10,10 +10,17 @@ import type { UserConfig } from "../common/config/userConfig.js";
1010
import type { Server } from "../server.js";
1111
import type { Elicitation } from "../elicitation.js";
1212
import type { PreviewFeature } from "../common/schemas.js";
13+
import { type ZodNever } from "zod";
1314

1415
export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNever>;
1516
export type ToolCallbackArgs<Args extends ZodRawShape> = Parameters<ToolCallback<Args>>;
1617

18+
export type ToolResult<OutputSchema extends ZodRawShape | undefined = undefined> = {
19+
content: { type: "text"; text: string }[];
20+
structuredContent: OutputSchema extends ZodRawShape ? z.objectOutputType<OutputSchema, ZodTypeAny> : never;
21+
isError?: boolean;
22+
};
23+
1724
export type ToolExecutionContext<Args extends ZodRawShape = ZodRawShape> = Parameters<ToolCallback<Args>>[1];
1825

1926
/**
@@ -274,6 +281,8 @@ export abstract class ToolBase {
274281
*/
275282
protected abstract argsShape: ZodRawShape;
276283

284+
protected outputShape?: ZodRawShape;
285+
277286
private registeredTool: RegisteredTool | undefined;
278287

279288
protected get annotations(): ToolAnnotations {
@@ -462,11 +471,14 @@ export abstract class ToolBase {
462471
}
463472
};
464473

465-
this.registeredTool = server.mcpServer.tool(
474+
this.registeredTool = server.mcpServer.registerTool(
466475
this.name,
467-
this.description,
468-
this.argsShape,
469-
this.annotations,
476+
{
477+
description: this.description,
478+
inputSchema: this.argsShape,
479+
annotations: this.annotations,
480+
outputSchema: this.outputShape,
481+
},
470482
callback
471483
);
472484

‎tests/integration/helpers.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ export function setupIntegrationTest(
184184
}
185185

186186
// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
187-
export function getResponseContent(content: unknown | { content: unknown }): string {
187+
export function getResponseContent(content: unknown | { content: unknown; structuredContent: unknown }): string {
188188
return getResponseElements(content)
189189
.map((item) => item.text)
190190
.join("\n");

‎tests/integration/tools/mongodb/create/insertMany.test.ts‎

Lines changed: 32 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -72,10 +72,11 @@ describeWithMongoDB("insertMany tool when search is disabled", (integration) =>
7272
},
7373
});
7474

75-
const content = getResponseContent(response.content);
75+
const content = getResponseContent(response);
7676
expect(content).toContain(`Inserted \`1\` document(s) into ${integration.randomDbName()}.coll1.`);
7777

7878
await validateDocuments("coll1", [{ prop1: "value1" }]);
79+
validateStructuredContent(response.structuredContent, extractInsertedIds(content));
7980
});
8081

8182
it("returns an error when inserting duplicates", async () => {
@@ -95,7 +96,7 @@ describeWithMongoDB("insertMany tool when search is disabled", (integration) =>
9596
},
9697
});
9798

98-
const content = getResponseContent(response.content);
99+
const content = getResponseContent(response);
99100
expect(content).toContain("Error running insert-many");
100101
expect(content).toContain("duplicate key error");
101102
expect(content).toContain(insertedIds[0]?.toString());
@@ -174,12 +175,14 @@ describeWithMongoDB(
174175
},
175176
});
176177

177-
const content = getResponseContent(response.content);
178+
const content = getResponseContent(response);
178179
const insertedIds = extractInsertedIds(content);
179180
expect(insertedIds).toHaveLength(1);
180181

181182
const docCount = await collection.countDocuments({ _id: insertedIds[0] });
182183
expect(docCount).toBe(1);
184+
185+
validateStructuredContent(response.structuredContent, insertedIds);
183186
});
184187

185188
it("returns an error when there is a search index and embeddings parameter are wrong", async () => {
@@ -214,7 +217,7 @@ describeWithMongoDB(
214217
},
215218
});
216219

217-
const content = getResponseContent(response.content);
220+
const content = getResponseContent(response);
218221
expect(content).toContain("Error running insert-many");
219222
const untrustedContent = getDataFromUntrustedContent(content);
220223
expect(untrustedContent).toContain(
@@ -263,10 +266,11 @@ describeWithMongoDB(
263266
},
264267
});
265268

266-
const content = getResponseContent(response.content);
269+
const content = getResponseContent(response);
267270
expect(content).toContain("Documents were inserted successfully.");
268271
const insertedIds = extractInsertedIds(content);
269272
expect(insertedIds).toHaveLength(1);
273+
validateStructuredContent(response.structuredContent, insertedIds);
270274

271275
const doc = await collection.findOne({ _id: insertedIds[0] });
272276
expect(doc).toBeDefined();
@@ -316,10 +320,11 @@ describeWithMongoDB(
316320
},
317321
});
318322

319-
const content = getResponseContent(response.content);
323+
const content = getResponseContent(response);
320324
expect(content).toContain("Documents were inserted successfully.");
321325
const insertedIds = extractInsertedIds(content);
322326
expect(insertedIds).toHaveLength(2);
327+
validateStructuredContent(response.structuredContent, insertedIds);
323328

324329
const doc1 = await collection.findOne({ _id: insertedIds[0] });
325330
expect(doc1?.title).toBe("The Matrix");
@@ -369,10 +374,11 @@ describeWithMongoDB(
369374
},
370375
});
371376

372-
const content = getResponseContent(response.content);
377+
const content = getResponseContent(response);
373378
expect(content).toContain("Documents were inserted successfully.");
374379
const insertedIds = extractInsertedIds(content);
375380
expect(insertedIds).toHaveLength(1);
381+
validateStructuredContent(response.structuredContent, insertedIds);
376382

377383
const doc = await collection.findOne({ _id: insertedIds[0] });
378384
expect(doc?.info).toBeDefined();
@@ -417,10 +423,11 @@ describeWithMongoDB(
417423
},
418424
});
419425

420-
const content = getResponseContent(response.content);
426+
const content = getResponseContent(response);
421427
expect(content).toContain("Documents were inserted successfully.");
422428
const insertedIds = extractInsertedIds(content);
423429
expect(insertedIds).toHaveLength(1);
430+
validateStructuredContent(response.structuredContent, insertedIds);
424431

425432
const doc = await collection.findOne({ _id: insertedIds[0] });
426433
expect(doc?.title).toBe("The Matrix");
@@ -452,10 +459,11 @@ describeWithMongoDB(
452459
},
453460
},
454461
});
455-
const content = getResponseContent(response.content);
462+
const content = getResponseContent(response);
456463
expect(content).toContain("Documents were inserted successfully.");
457464
const insertedIds = extractInsertedIds(content);
458465
expect(insertedIds).toHaveLength(1);
466+
validateStructuredContent(response.structuredContent, insertedIds);
459467

460468
const doc = await collection.findOne({ _id: insertedIds[0] });
461469
expect((doc?.title as Record<string, unknown>)?.text).toBe("The Matrix");
@@ -495,7 +503,7 @@ describeWithMongoDB(
495503
},
496504
});
497505

498-
const content = getResponseContent(response.content);
506+
const content = getResponseContent(response);
499507
expect(content).toContain("Error running insert-many");
500508
expect(content).toContain("Field 'nonExistentField' does not have a vector search index in collection");
501509
expect(content).toContain("Only fields with vector search indexes can have embeddings generated");
@@ -529,10 +537,11 @@ describeWithMongoDB(
529537
},
530538
});
531539

532-
const content = getResponseContent(response.content);
540+
const content = getResponseContent(response);
533541
expect(content).toContain("Documents were inserted successfully.");
534542
const insertedIds = extractInsertedIds(content);
535543
expect(insertedIds).toHaveLength(1);
544+
validateStructuredContent(response.structuredContent, insertedIds);
536545

537546
const doc = await collection.findOne({ _id: insertedIds[0] });
538547
expect(doc?.title).toBe("The Matrix");
@@ -564,9 +573,10 @@ describeWithMongoDB(
564573
},
565574
});
566575

567-
const content = getResponseContent(response.content);
576+
const content = getResponseContent(response);
568577
expect(content).toContain("Documents were inserted successfully.");
569578
const insertedIds = extractInsertedIds(content);
579+
validateStructuredContent(response.structuredContent, insertedIds);
570580

571581
const doc = await collection.findOne({ _id: insertedIds[0] });
572582
expect(Array.isArray(doc?.titleEmbeddings)).toBe(true);
@@ -614,9 +624,10 @@ describeWithMongoDB(
614624
},
615625
});
616626

617-
const content = getResponseContent(response.content);
627+
const content = getResponseContent(response);
618628
expect(content).toContain("Documents were inserted successfully.");
619629
const insertedIds = extractInsertedIds(content);
630+
validateStructuredContent(response.structuredContent, insertedIds);
620631

621632
const doc = await collection.findOne({ _id: insertedIds[0] });
622633
expect(doc?.title).toBe("The Matrix");
@@ -692,3 +703,11 @@ function extractInsertedIds(content: string): ObjectId[] {
692703
.map((e) => ObjectId.createFromHexString(e)) ?? []
693704
);
694705
}
706+
707+
function validateStructuredContent(structuredContent: unknown, expectedIds: ObjectId[]): void {
708+
expect(structuredContent).toEqual({
709+
success: true,
710+
insertedCount: expectedIds.length,
711+
insertedIds: expectedIds,
712+
});
713+
}

0 commit comments

Comments
 (0)