Skip to content

Commit 2e4e049

Browse files
kmruizkmruiz
authored
Merge branch 'main' into chore/mcp-245 
2 parents 9fd4634 + f35f0e5 commit 2e4e049

19 files changed

+1310
-661
lines changed

‎.github/workflows/cleanup-atlas-env.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ name: "Cleanup stale Atlas test environments"
33
on:
44
workflow_dispatch:
55
schedule:
6-
- cron: "0 0 * * *"
6+
- cron: "0 * * * *"
77

88
permissions: {}
99

‎src/server.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import {
1818
UnsubscribeRequestSchema,
1919
} from "@modelcontextprotocol/sdk/types.js";
2020
import assert from "assert";
21-
import type { ToolBase, ToolConstructorParams } from "./tools/tool.js";
21+
import type { ToolBase, ToolCategory, ToolConstructorParams } from "./tools/tool.js";
2222
import { validateConnectionString } from "./helpers/connectionOptions.js";
2323
import { packageInfo } from "./common/packageInfo.js";
2424
import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
@@ -174,6 +174,10 @@ export class Server {
174174
this.mcpServer.sendResourceListChanged();
175175
}
176176

177+
public isToolCategoryAvailable(name: ToolCategory): boolean {
178+
return !!this.tools.filter((t) => t.category === name).length;
179+
}
180+
177181
public sendResourceUpdated(uri: string): void {
178182
this.session.logger.info({
179183
id: LogId.resourceUpdateFailure,

‎src/tools/mongodb/delete/dropIndex.ts‎

Lines changed: 60 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,42 @@
11
import z from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3+
import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
34
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
4-
import { type ToolArgs, type OperationType, formatUntrustedData } from "../../tool.js";
5+
import { type ToolArgs, type OperationType, formatUntrustedData, FeatureFlags } from "../../tool.js";
56

67
export class DropIndexTool extends MongoDBToolBase {
78
public name = "drop-index";
89
protected description = "Drop an index for the provided database and collection.";
910
protected argsShape = {
1011
...DbOperationArgs,
1112
indexName: z.string().nonempty().describe("The name of the index to be dropped."),
13+
type: this.isFeatureFlagEnabled(FeatureFlags.VectorSearch)
14+
? z
15+
.enum(["classic", "search"])
16+
.describe(
17+
"The type of index to be deleted. Use 'classic' for standard indexes and 'search' for atlas search and vector search indexes."
18+
)
19+
: z
20+
.literal("classic")
21+
.default("classic")
22+
.describe("The type of index to be deleted. Is always set to 'classic'."),
1223
};
1324
public operationType: OperationType = "delete";
1425

15-
protected async execute({
16-
database,
17-
collection,
18-
indexName,
19-
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
26+
protected async execute(toolArgs: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
2027
const provider = await this.ensureConnected();
28+
switch (toolArgs.type) {
29+
case "classic":
30+
return this.dropClassicIndex(provider, toolArgs);
31+
case "search":
32+
return this.dropSearchIndex(provider, toolArgs);
33+
}
34+
}
35+
36+
private async dropClassicIndex(
37+
provider: NodeDriverServiceProvider,
38+
{ database, collection, indexName }: ToolArgs<typeof this.argsShape>
39+
): Promise<CallToolResult> {
2140
const result = await provider.runCommand(database, {
2241
dropIndexes: collection,
2342
index: indexName,
@@ -35,9 +54,42 @@ export class DropIndexTool extends MongoDBToolBase {
3554
};
3655
}
3756

38-
protected getConfirmationMessage({ database, collection, indexName }: ToolArgs<typeof this.argsShape>): string {
57+
private async dropSearchIndex(
58+
provider: NodeDriverServiceProvider,
59+
{ database, collection, indexName }: ToolArgs<typeof this.argsShape>
60+
): Promise<CallToolResult> {
61+
await this.ensureSearchIsSupported();
62+
const indexes = await provider.getSearchIndexes(database, collection, indexName);
63+
if (indexes.length === 0) {
64+
return {
65+
content: formatUntrustedData(
66+
"Index does not exist in the provided namespace.",
67+
JSON.stringify({ indexName, namespace: `${database}.${collection}` })
68+
),
69+
isError: true,
70+
};
71+
}
72+
73+
await provider.dropSearchIndex(database, collection, indexName);
74+
return {
75+
content: formatUntrustedData(
76+
"Successfully dropped the index from the provided namespace.",
77+
JSON.stringify({
78+
indexName,
79+
namespace: `${database}.${collection}`,
80+
})
81+
),
82+
};
83+
}
84+
85+
protected getConfirmationMessage({
86+
database,
87+
collection,
88+
indexName,
89+
type,
90+
}: ToolArgs<typeof this.argsShape>): string {
3991
return (
40-
`You are about to drop the \`${indexName}\` index from the \`${database}.${collection}\` namespace:\n\n` +
92+
`You are about to drop the ${type === "search" ? "search index" : "index"} named \`${indexName}\` from the \`${database}.${collection}\` namespace:\n\n` +
4193
"This operation will permanently remove the index and might affect the performance of queries relying on this index.\n\n" +
4294
"**Do you confirm the execution of the action?**"
4395
);

‎src/tools/mongodb/metadata/collectionIndexes.ts‎

Lines changed: 52 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,20 @@
11
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
33
import type { ToolArgs, OperationType } from "../../tool.js";
4-
import { formatUntrustedData } from "../../tool.js";
4+
import { FeatureFlags, formatUntrustedData } from "../../tool.js";
5+
6+
type SearchIndexStatus = {
7+
name: string;
8+
type: string;
9+
status: string;
10+
queryable: boolean;
11+
latestDefinition: Document;
12+
};
13+
14+
type IndexStatus = {
15+
name: string;
16+
key: Document;
17+
};
518

619
export class CollectionIndexesTool extends MongoDBToolBase {
720
public name = "collection-indexes";
@@ -12,12 +25,30 @@ export class CollectionIndexesTool extends MongoDBToolBase {
1225
protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
1326
const provider = await this.ensureConnected();
1427
const indexes = await provider.getIndexes(database, collection);
28+
const indexDefinitions: IndexStatus[] = indexes.map((index) => ({
29+
name: index.name as string,
30+
key: index.key as Document,
31+
}));
32+
33+
const searchIndexDefinitions: SearchIndexStatus[] = [];
34+
if (this.isFeatureFlagEnabled(FeatureFlags.VectorSearch) && (await this.session.isSearchSupported())) {
35+
const searchIndexes = await provider.getSearchIndexes(database, collection);
36+
searchIndexDefinitions.push(...this.extractSearchIndexDetails(searchIndexes));
37+
}
1538

1639
return {
17-
content: formatUntrustedData(
18-
`Found ${indexes.length} indexes in the collection "${collection}":`,
19-
...indexes.map((index) => `Name: "${index.name}", definition: ${JSON.stringify(index.key)}`)
20-
),
40+
content: [
41+
...formatUntrustedData(
42+
`Found ${indexDefinitions.length} indexes in the collection "${collection}":`,
43+
...indexDefinitions.map((i) => JSON.stringify(i))
44+
),
45+
...(searchIndexDefinitions.length > 0
46+
? formatUntrustedData(
47+
`Found ${searchIndexDefinitions.length} search and vector search indexes in the collection "${collection}":`,
48+
...searchIndexDefinitions.map((i) => JSON.stringify(i))
49+
)
50+
: []),
51+
],
2152
};
2253
}
2354

@@ -39,4 +70,20 @@ export class CollectionIndexesTool extends MongoDBToolBase {
3970

4071
return super.handleError(error, args);
4172
}
73+
74+
/**
75+
* Atlas Search index status contains a lot of information that is not relevant for the agent at this stage.
76+
* Like for example, the status on each of the dedicated nodes. We only care about the main status, if it's
77+
* queryable and the index name. We are also picking the index definition as it can be used by the agent to
78+
* understand which fields are available for searching.
79+
**/
80+
protected extractSearchIndexDetails(indexes: Record<string, unknown>[]): SearchIndexStatus[] {
81+
return indexes.map((index) => ({
82+
name: (index["name"] ?? "default") as string,
83+
type: (index["type"] ?? "UNKNOWN") as string,
84+
status: (index["status"] ?? "UNKNOWN") as string,
85+
queryable: (index["queryable"] ?? false) as boolean,
86+
latestDefinition: index["latestDefinition"] as Document,
87+
}));
88+
}
4289
}

‎src/tools/mongodb/mongodbTool.ts‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ export abstract class MongoDBToolBase extends ToolBase {
8787
isError: true,
8888
};
8989
case ErrorCodes.AtlasSearchNotSupported: {
90-
const CTA = this.isToolCategoryAvailable("atlas-local" as unknown as ToolCategory)
90+
const CTA = this.server?.isToolCategoryAvailable("atlas-local" as unknown as ToolCategory)
9191
? "`atlas-local` tools"
9292
: "Atlas CLI";
9393
return {
@@ -123,8 +123,4 @@ export abstract class MongoDBToolBase extends ToolBase {
123123

124124
return metadata;
125125
}
126-
127-
protected isToolCategoryAvailable(name: ToolCategory): boolean {
128-
return (this.server?.tools.filter((t) => t.category === name).length ?? 0) > 0;
129-
}
130126
}

‎src/tools/mongodb/search/listSearchIndexes.ts‎

Lines changed: 0 additions & 65 deletions
This file was deleted.

‎src/tools/mongodb/tools.ts‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ import { ExplainTool } from "./metadata/explain.js";
1919
import { CreateCollectionTool } from "./create/createCollection.js";
2020
import { LogsTool } from "./metadata/logs.js";
2121
import { ExportTool } from "./read/export.js";
22-
import { ListSearchIndexesTool } from "./search/listSearchIndexes.js";
2322
import { DropIndexTool } from "./delete/dropIndex.js";
2423

2524
export const MongoDbTools = [
@@ -45,5 +44,4 @@ export const MongoDbTools = [
4544
CreateCollectionTool,
4645
LogsTool,
4746
ExportTool,
48-
ListSearchIndexesTool,
4947
];

‎tests/accuracy/collectionIndexes.test.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,4 +37,28 @@ describeAccuracyTests([
3737
},
3838
],
3939
},
40+
{
41+
prompt: "how many search indexes do I have in the collection mydb.mycoll?",
42+
expectedToolCalls: [
43+
{
44+
toolName: "collection-indexes",
45+
parameters: {
46+
database: "mydb",
47+
collection: "mycoll",
48+
},
49+
},
50+
],
51+
},
52+
{
53+
prompt: "which vector search indexes do I have in mydb.mycoll?",
54+
expectedToolCalls: [
55+
{
56+
toolName: "collection-indexes",
57+
parameters: {
58+
database: "mydb",
59+
collection: "mycoll",
60+
},
61+
},
62+
],
63+
},
4064
]);

0 commit comments

Comments
 (0)