Skip to content

Commit 7cc4f65

Browse files
duwenxin99averikitsch
duwenxin99
and
averikitsch
authored
feat(server): add support for toolset filtering in prebuilt CLI flag (googleapis#3245)
Toolset Filtering for Prebuilt Configs: Enabled loading specific toolsets from a prebuilt configuration by appending a / suffix (e.g., --prebuilt sqlite/sqlite_database_tools), allowing users to load only a subset of tools rather than the entire configuration. --------- Co-authored-by: Averi Kitsch <akitsch@google.com>
1 parent 7ed1e7b commit 7cc4f65

6 files changed

Lines changed: 199 additions & 6 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ You can use the Toolbox in any MCP-compatible IDE or client (e.g., Gemini CLI, G
9595

9696
2. Set the appropriate environment variables to connect, see the [Prebuilt Tools Reference](https://mcp-toolbox.dev/documentation/configuration/prebuilt-configs/).
9797

98-
When you run Toolbox with a `--prebuilt=<database>` flag, you instantly get access to standard tools to interact with that database.
98+
When you run Toolbox with a `--prebuilt=<database>` flag, you instantly get access to standard tools to interact with that database. You can also specify a specific toolset using the `--prebuilt=<database>/<toolset>` syntax (e.g., `--prebuilt=postgres/data` to only load SQL tools).
9999

100100
Supported databases currently include:
101101
- **Google Cloud:** AlloyDB, BigQuery, Cloud SQL (PostgreSQL, MySQL, SQL Server), Spanner, Firestore, Knowledge Catalog (formerly known as Dataplex).

‎cmd/internal/options.go‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,22 @@ func (opts *ToolboxOptions) LoadConfig(ctx context.Context, parser *ConfigParser
206206
logger.InfoContext(ctx, logMsg)
207207

208208
for _, configName := range opts.PrebuiltConfigs {
209-
buf, err := prebuiltconfigs.Get(configName)
209+
if !strings.Contains(configName, "/") {
210+
for _, sep := range []string{".", ":", "@"} {
211+
if strings.Contains(configName, sep) {
212+
parts := strings.SplitN(configName, sep, 2)
213+
if slices.Contains(prebuiltconfigs.GetPrebuiltSources(), parts[0]) {
214+
errMsg := fmt.Errorf("invalid prebuilt config format '%s'. Did you mean '%s/%s'? Use '/' to specify a toolset", configName, parts[0], parts[1])
215+
logger.ErrorContext(ctx, errMsg.Error())
216+
return isCustomConfigured, errMsg
217+
}
218+
}
219+
}
220+
}
221+
222+
sourceName, toolsetName, _ := strings.Cut(configName, "/")
223+
224+
buf, err := prebuiltconfigs.Get(sourceName)
210225
if err != nil {
211226
logger.ErrorContext(ctx, err.Error())
212227
return isCustomConfigured, err
@@ -219,6 +234,35 @@ func (opts *ToolboxOptions) LoadConfig(ctx context.Context, parser *ConfigParser
219234
logger.ErrorContext(ctx, errMsg.Error())
220235
return isCustomConfigured, errMsg
221236
}
237+
238+
if toolsetName != "" {
239+
targetToolset, exists := parsed.Toolsets[toolsetName]
240+
if !exists {
241+
var available []string
242+
for k := range parsed.Toolsets {
243+
available = append(available, k)
244+
}
245+
slices.Sort(available)
246+
errMsg := fmt.Errorf("toolset '%s' not found in prebuilt configuration '%s'. Available toolsets: %s", toolsetName, sourceName, strings.Join(available, ", "))
247+
logger.ErrorContext(ctx, errMsg.Error())
248+
return isCustomConfigured, errMsg
249+
}
250+
251+
// Filter tools to only include those in the target toolset
252+
filteredTools := make(server.ToolConfigs)
253+
for _, tName := range targetToolset.ToolNames {
254+
if tCfg, tExists := parsed.Tools[tName]; tExists {
255+
filteredTools[tName] = tCfg
256+
}
257+
}
258+
parsed.Tools = filteredTools
259+
260+
// Filter toolsets to only include the target toolset
261+
filteredToolsets := make(server.ToolsetConfigs)
262+
filteredToolsets[toolsetName] = targetToolset
263+
parsed.Toolsets = filteredToolsets
264+
}
265+
222266
allConfigs = append(allConfigs, parsed)
223267
}
224268
}
@@ -241,7 +285,8 @@ func (opts *ToolboxOptions) LoadConfig(ctx context.Context, parser *ConfigParser
241285
}
242286
// prebuiltConfigs is already sorted above
243287
for _, configName := range opts.PrebuiltConfigs {
244-
opts.Cfg.Version += fmt.Sprintf("+%s.%s", tag, configName)
288+
sanitizedConfigName := strings.ReplaceAll(configName, "/", ".")
289+
opts.Cfg.Version += fmt.Sprintf("+%s.%s", tag, sanitizedConfigName)
245290
}
246291
}
247292

‎cmd/internal/options_test.go‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,11 @@ package internal
1717
import (
1818
"errors"
1919
"io"
20+
"strings"
2021
"testing"
22+
23+
"github.com/googleapis/mcp-toolbox/internal/server"
24+
"github.com/googleapis/mcp-toolbox/internal/testutils"
2125
)
2226

2327
func TestToolboxOptions(t *testing.T) {
@@ -47,3 +51,100 @@ func TestToolboxOptions(t *testing.T) {
4751
})
4852
}
4953
}
54+
55+
func TestLoadConfig(t *testing.T) {
56+
t.Setenv("CLOUD_HEALTHCARE_PROJECT", "mock")
57+
t.Setenv("CLOUD_HEALTHCARE_REGION", "mock")
58+
t.Setenv("CLOUD_HEALTHCARE_DATASET", "mock")
59+
t.Setenv("BIGQUERY_PROJECT", "mock")
60+
t.Setenv("POSTGRES_HOST", "localhost")
61+
t.Setenv("POSTGRES_PORT", "5432")
62+
t.Setenv("POSTGRES_DATABASE", "mock")
63+
t.Setenv("POSTGRES_USER", "mock")
64+
t.Setenv("POSTGRES_PASSWORD", "mock")
65+
66+
ctx, err := testutils.ContextWithNewLogger()
67+
if err != nil {
68+
t.Fatalf("unexpected error: %s", err)
69+
}
70+
71+
tcs := []struct {
72+
desc string
73+
prebuiltConfigs []string
74+
initialVersion string
75+
wantVersion string
76+
wantErr string
77+
matchPrefix bool
78+
}{
79+
{
80+
desc: "version sanitization with multiple configs",
81+
prebuiltConfigs: []string{"cloud-healthcare/cloud_healthcare_fhir_tools", "bigquery"},
82+
initialVersion: "v1.0.0",
83+
wantVersion: "v1.0.0+prebuilt.bigquery+prebuilt.cloud-healthcare.cloud_healthcare_fhir_tools",
84+
},
85+
{
86+
desc: "toolset not found in prebuilt config",
87+
prebuiltConfigs: []string{"postgres/invalid-toolset"},
88+
wantErr: "toolset 'invalid-toolset' not found in prebuilt configuration 'postgres'. Available toolsets: data, health, monitor, replication, view-config",
89+
},
90+
{
91+
desc: "invalid separator - dot",
92+
prebuiltConfigs: []string{"postgres.sql"},
93+
wantErr: "invalid prebuilt config format 'postgres.sql'. Did you mean 'postgres/sql'? Use '/' to specify a toolset",
94+
},
95+
{
96+
desc: "invalid separator - colon",
97+
prebuiltConfigs: []string{"postgres:sql"},
98+
wantErr: "invalid prebuilt config format 'postgres:sql'. Did you mean 'postgres/sql'? Use '/' to specify a toolset",
99+
},
100+
{
101+
desc: "invalid separator - at",
102+
prebuiltConfigs: []string{"postgres@sql"},
103+
wantErr: "invalid prebuilt config format 'postgres@sql'. Did you mean 'postgres/sql'? Use '/' to specify a toolset",
104+
},
105+
{
106+
desc: "no warning on unrelated dots",
107+
prebuiltConfigs: []string{"invalid-source.sql"},
108+
wantErr: "prebuilt source tool for 'invalid-source.sql' not found",
109+
matchPrefix: true,
110+
},
111+
}
112+
113+
for _, tc := range tcs {
114+
t.Run(tc.desc, func(t *testing.T) {
115+
opts := &ToolboxOptions{
116+
PrebuiltConfigs: tc.prebuiltConfigs,
117+
Cfg: server.ServerConfig{
118+
Version: tc.initialVersion,
119+
},
120+
}
121+
122+
parser := &ConfigParser{}
123+
_, err = opts.LoadConfig(ctx, parser)
124+
125+
if tc.wantVersion != "" {
126+
// Success case
127+
if err != nil {
128+
t.Fatalf("unexpected error loading config: %v", err)
129+
}
130+
if opts.Cfg.Version != tc.wantVersion {
131+
t.Errorf("unexpected version: got %q, want %q", opts.Cfg.Version, tc.wantVersion)
132+
}
133+
} else {
134+
// Failure case
135+
if err == nil {
136+
t.Fatalf("expected error, got nil")
137+
}
138+
if tc.matchPrefix {
139+
if !strings.HasPrefix(err.Error(), tc.wantErr) {
140+
t.Errorf("unexpected error message:\ngot: %q\nwant prefix: %q", err.Error(), tc.wantErr)
141+
}
142+
} else {
143+
if err.Error() != tc.wantErr {
144+
t.Errorf("unexpected error message:\ngot: %q\nwant: %q", err.Error(), tc.wantErr)
145+
}
146+
}
147+
}
148+
})
149+
}
150+
}

‎cmd/root_test.go‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -412,6 +412,11 @@ func TestPrebuiltFlag(t *testing.T) {
412412
args: []string{"--prebuilt", "alloydb,bigquery"},
413413
want: []string{"alloydb", "bigquery"},
414414
},
415+
{
416+
desc: "prebuilt toolset flag",
417+
args: []string{"--prebuilt", "alloydb-postgres/monitor"},
418+
want: []string{"alloydb-postgres/monitor"},
419+
},
415420
}
416421
for _, tc := range tcs {
417422
t.Run(tc.desc, func(t *testing.T) {
@@ -857,6 +862,42 @@ tools:
857862
wantErr: true,
858863
errString: "resource conflicts detected",
859864
},
865+
{
866+
desc: "success toolset filtering",
867+
args: []string{"--prebuilt", "sqlite/sqlite_database_tools"},
868+
wantErr: false,
869+
cfgCheck: func(cfg server.ServerConfig) error {
870+
if _, ok := cfg.ToolConfigs["execute_sql"]; !ok {
871+
return fmt.Errorf("expected tool 'execute_sql' not found")
872+
}
873+
if _, ok := cfg.ToolConfigs["list_tables"]; !ok {
874+
return fmt.Errorf("expected tool 'list_tables' not found")
875+
}
876+
if len(cfg.ToolConfigs) != 2 {
877+
return fmt.Errorf("expected exactly 2 tools, got %d", len(cfg.ToolConfigs))
878+
}
879+
if _, ok := cfg.ToolsetConfigs["sqlite_database_tools"]; !ok {
880+
return fmt.Errorf("expected toolset 'sqlite_database_tools' not found")
881+
}
882+
if len(cfg.ToolsetConfigs) != 2 {
883+
var names []string
884+
for k := range cfg.ToolsetConfigs {
885+
names = append(names, k)
886+
}
887+
return fmt.Errorf("expected exactly 2 toolsets (including default), got %d: %v", len(cfg.ToolsetConfigs), names)
888+
}
889+
if _, ok := cfg.ToolsetConfigs[""]; !ok {
890+
return fmt.Errorf("expected default toolset '' not found")
891+
}
892+
return nil
893+
},
894+
},
895+
{
896+
desc: "toolset not found error",
897+
args: []string{"--prebuilt", "sqlite/nonexistent"},
898+
wantErr: true,
899+
errString: "toolset 'nonexistent' not found in prebuilt configuration 'sqlite'",
900+
},
860901
}
861902

862903
for _, tc := range testCases {
@@ -882,7 +923,7 @@ tools:
882923
}
883924
if tc.cfgCheck != nil {
884925
if err := tc.cfgCheck(opts.Cfg); err != nil {
885-
t.Errorf("config check failed: %v", err)
926+
t.Errorf("config check failed: %v. Output:\n%s", err, output)
886927
}
887928
}
888929
}

‎docs/en/documentation/configuration/prebuilt-configs/_index.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@ You can now use `--prebuilt` along `--config`, `--configs`, or
1919

2020
You can also combine multiple prebuilt configs.
2121

22+
**Filtering Toolsets:**
23+
You can load a specific toolset from a prebuilt configuration by appending a `/` and the toolset name, for example: `--prebuilt=postgres/data` to only load the SQL tools.
24+
2225
See [Usage Examples](../../../reference/cli.md#usage-examples).
2326
{{< /notice >}}
2427

‎docs/en/reference/cli.md‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ description: >
2121
| `-p` | `--port` | Port the server will listen on. | `5000` |
2222
| | `--tls-cert` | Path to the PEM-encoded TLS certificate file. | |
2323
| | `--tls-key` | Path to the PEM-encoded TLS private key file. | |
24-
| | `--prebuilt` | Use one or more prebuilt tool configuration by source type. See [Prebuilt Tools Reference](../documentation/configuration/prebuilt-configs/_index.md) for allowed values. | |
24+
| | `--prebuilt` | Use one or more prebuilt tool configuration by source type. Optionally specify a toolset suffix (e.g., `<source>/<toolset>`) to load only that toolset. See [Prebuilt Tools Reference](../documentation/configuration/prebuilt-configs/_index.md) for allowed values. | |
2525
| | `--stdio` | Listens via MCP STDIO instead of acting as a remote HTTP server. | |
2626
| | `--telemetry-gcp` | Enable exporting directly to Google Cloud Monitoring. | |
2727
| | `--telemetry-gcp-project` | Google Cloud project ID used for `--telemetry-gcp`; defaults to `GOOGLE_CLOUD_PROJECT` if not set. | |
@@ -161,6 +161,9 @@ By default, traffic is unencrypted (HTTP). In production or shared networks, you
161161
./toolbox --prebuilt alloydb-postgres,alloydb-postgres-admin
162162
# OR
163163
./toolbox --prebuilt alloydb-postgres --prebuilt alloydb-postgres-admin
164+
165+
# Server filtering a prebuilt configuration to load only a specific toolset
166+
./toolbox --prebuilt alloydb-postgres/monitor
164167
```
165168

166169
### Tool Configuration Sources
@@ -182,7 +185,7 @@ The CLI supports multiple mutually exclusive ways to specify tool configurations
182185
**Prebuilt Configurations:**
183186

184187
- `--prebuilt`: Use one or more predefined configurations for specific database types (e.g.,
185-
'bigquery', 'postgres', 'spanner'). See [Prebuilt Tools
188+
'bigquery', 'postgres', 'spanner'), optionally appending a toolset name to filter the loaded tools (e.g., `alloydb-postgres/monitor`). See [Prebuilt Tools
186189
Reference](../documentation/configuration/prebuilt-configs/_index.md) for allowed values.
187190

188191
{{< notice tip >}}

0 commit comments

Comments
 (0)