[SWITCHin2]

Description

Adds source-level collection scoping for MongoDB, and shows the allowed set to agents as a JSON Schema enum.

• allowedCollections on the mongodb source. This optional list of database.collection entries applies to every tool on the source, the same way BigQuery's allowedDatasets does. Leaving it unset keeps today's behaviour (no restriction).
• Source is authoritative, the tool can only narrow it. A tool's collectionAllowedValues narrows the source's list. The effective set is source ∩ tool.
• Checked at startup. A tool whose fixed collection is outside the source list, or whose effective set is empty, fails when the server starts.
• Checked at invoke. Each MongoDB tool checks the resolved collection again with IsCollectionAllowed. If it's not allowed, the agent gets an error that lists the allowed collections, so it can retry.
• Aggregate pipelines are checked too. Collections named by $lookup, $graphLookup, $unionWith, $out, $merge and $facet stages (including nested pipelines) must also be allowed. Without this, a pipeline could reach collections outside the scope.
• JSON Schema enum in MCP manifests. A parameter whose allowedValues are all literal strings now shows them as enum in tools/list. This lives in the shared parameters package, so every tool that uses allowedValues gets it. Regex entries and non-string values produce no enum, because some LLM function-calling schemas only accept string enums.

Precedence

source.allowedCollections	tool config	behaviour
unset	fixed collection: c	allowed, as today
unset	runtime, no collectionAllowedValues	runtime param, no enum
unset	runtime + collectionAllowedValues: [a,b]	runtime param, enum: [a,b]
[a,b,c]	fixed collection: a	allowed
[a,b,c]	fixed collection: z	rejected at startup
[a,b,c]	runtime	runtime param, enum: [a,b,c]
[a,b,c]	runtime + collectionAllowedValues: [a,b]	runtime param, enum: [a,b]
[a,b]	runtime + collectionAllowedValues: [c]	rejected at startup

Backward compatibility

• allowedCollections is optional. When it's unset, every check returns early and existing configs behave the same as before.
• The existing rule that collection and collectionAllowedValues can't both be set is unchanged.
• enum has omitempty, and only the MCP manifest changes; the native Toolbox manifest is untouched. It only makes visible what the server already enforced through allowedValues.

Follow-up: the mongodb-list-collections tool is in a stacked PR on top of this one.

PR Checklist

• Make sure you reviewed
  CONTRIBUTING.md
• Make sure to open an issue as a
  bug/issue
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
• Ensure you have manually reviewed the entire diff before requesting a
  review
• Ensure the tests and linter pass
• Code coverage does not decrease (if any source code was changed)
• Appropriate docs were updated (if necessary)
• Make sure to add ! if this involve a breaking change

[gemini-code-assist]

Code Review

This pull request implements collection scoping for MongoDB sources and tools, allowing administrators to restrict database access to a specified list of collections. It updates MongoDB tool configurations, validation logic, and runtime parameter resolution to enforce these restrictions, and advertises literal allowed values as JSON Schema enums to MCP clients. The feedback suggests enhancing the robustness of the 'asDocument' and 'asArray' helper functions in 'internal/sources/mongodb/mongodb.go' by using type assertions to support standard Go map and slice types alongside BSON-specific types.

[gemini-code-assist]

To make asDocument more robust and support standard map[string]any alongside bson.M and bson.D without compilation errors (regardless of whether bson.M is a type alias or a distinct type), use type assertions instead of a type switch.

func asDocument(v any) map[string]any {
    if m, ok := v.(map[string]any); ok {
        return m
    }
    if m, ok := v.(bson.M); ok {
        return m
    }
    if d, ok := v.(bson.D); ok {
        m := make(map[string]any, len(d))
        for _, e := range d {
            m[e.Key] = e.Value
        }
        return m
    }
    return nil
}

[SWITCHin2]

asDocument only ever gets values straight from bson.UnmarshalExtJSON: bson.M for the top-level stages and bson.D for nested documents. I checked that before writing the walker. A plain map[string]any never comes out of that decode, so the extra branch would never run.

[gemini-code-assist]

To make asArray more robust and support standard []any alongside bson.A without compilation errors (regardless of whether bson.A is a type alias or a distinct type), use type assertions for both types.

func asArray(v any) []any {
    if a, ok := v.([]any); ok {
        return a
    }
    if a, ok := v.(bson.A); ok {
        return a
    }
    return nil
}

[SWITCHin2]

Same reasoning as the asDocument thread: nested arrays from bson.UnmarshalExtJSON always decode as bson.A, so a plain []any can't reach this helper.

[AlexTalreja]

Using the regex here might be problematic if there are overlapping collection names. For example, if source allows orders and orders_confidential and the tool tries to restrict it to only orders, it seems like EffectiveCollections would still add both due to regex match, invalidating the tool restriction. What do you think?

[SWITCHin2]

My bad, that is something I implented wrong, Each collectionAllowedValues entry must now match the whole collection name, so orders no longer selects orders_confidential. Invoke also checks the effective set with an exact match, because the parameter's allowedValues check doesn't use anchors

[AlexTalreja]

Potential security issue: If bson.D contains duplicate keys, asDocument would overwrite the earlier key with the latest key before the checkPipelineScope check, in this case only passing allowed_collection through the check.

"$lookup": {
      "from": "secret_collection",
      "from": "allowed_collection",
      "localField": "_id",
       ...
}

However, after passing the check it then still sends the original bson.D (containing secret_collection) in Collection.Aggregate().

[SWITCHin2]

agreed. The scope check now rejects any pipeline document with a repeated key. It also rejects stages it can't read, such as a malformed from/into or a pipeline that isn't an array. The original bson.D is still what gets sent, so key order is kept.

Do let me know if this fixes the security issue, or some uniform way we follow on this project to prevent things like this.

[AlexTalreja]

In cases where the tool restricts collections to a narrower set than source, would this error message be incorrect since it's appending the entire source.MongoDBAllowedCollections(t.Cfg.Database))?

If so please update this (and the Invoke()) method in the other MongoDB tools so it doesn't provide misleading instructions to the agent.

[SWITCHin2]

Makes sense
All MongoDB tools now use one shared CheckCollectionAllowed helper, and the error lists only the effective (source ∩ tool) set.

[AlexTalreja]

This might not be a sufficient check, since allowedValues on an integer (or boolean) parameter can be represented as an array of strings, for example:

parameters:
  - name: page_count
    type: integer
    description: Number of results
    allowedValues: ["10", "25", "50"]

This passes v.(string) but then causes Toolbox to send an incorrect JSON schema with mismatched types to the LLM.

[SWITCHin2]

Fixed that, the new behavior is the enum is now only sent for type: string parameters, so integer and boolean parameters behave as they do on main.

[AlexTalreja]

/gcbrun

[AlexTalreja]

If a stage references a collection in another database via { "db": "other_db", "coll": "other_coll" } and also includes a nested pipeline, from what I understand it seems like appendCollectionRef resolves db locally for that reference, but subPipelineCollections still passes the outer database into pipelineCollections().

As a result, any collection references inside the sub-pipeline will be checked against allowedCollections for the outer database rather than other_db.

We should probably pass the resolved database from the parent stage into the recursive call or cover this case another way.

[SWITCHin2]

The nested pipeline is now checked in the database resolved from the stage's own from/coll, not the outer one. Added tests for both directions