From e535b372ea81864d644a67135a1b07e4e519b4b4 Mon Sep 17 00:00:00 2001 From: dishaprakash <57954147+dishaprakash@users.noreply.github.com> Date: Thu, 22 Jan 2026 23:00:17 +0000 Subject: [PATCH] feat: Support combining multiple prebuilt configurations (#2295) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description This PR introduces support for merging multiple prebuilt configurations. To ensure compatibility, the following restrictions apply: - No Naming Collisions: Configurations cannot share duplicate names for any resources (Tools, Sources, Toolsets, Auth Services, etc.). - Shared Environment Variables: If multiple sources rely on the same environment variable, they must share the same value; unique values for the same variable are not supported ## Usage Examples ### Successful Initialization You can load multiple prebuilt configurations by either repeating the --prebuilt flag or by providing a comma-separated list. **Option 1:** Multiple Flags ``` ./toolbox --prebuilt alloydb-postgres --prebuilt alloydb-postgres-admin ``` **Option 2:** Comma-Separated Values ``` ./toolbox --prebuilt alloydb-postgres,alloydb-postgres-admin ``` ### Initialization Failure (Resource Conflict) If two or more configurations define a resource with the same name (such as a Tool or Source, etc.), the server will fail to start and display a conflict error. ``` ./toolbox --prebuilt alloydb-postgres --prebuilt cloud-sql-mysql 2026-01-13T11:14:50.758121799Z INFO "Using prebuilt tool configurations for: alloydb-postgres, cloud-sql-mysql" 2026-01-13T11:14:50.764578167Z ERROR "resource conflicts detected:\n - tool 'execute_sql' (file #2)\n - tool 'list_active_queries' (file #2)\n - tool 'get_query_plan' (file #2)\n - tool 'list_tables' (file #2)\n\nPlease ensure each source, authService, tool, toolset and prompt has a unique name across all files" ``` ## PR Checklist > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [x] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [x] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #1855 --------- Co-authored-by: Averi Kitsch --- cmd/root.go | 64 ++++++++++++++++------------- cmd/root_test.go | 30 ++++++++++---- docs/en/reference/cli.md | 9 +++- docs/en/reference/prebuilt-tools.md | 3 ++ 4 files changed, 69 insertions(+), 37 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 4c373bf181f..dfac1c250f2 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -315,15 +315,15 @@ func Execute() { type Command struct { *cobra.Command - cfg server.ServerConfig - logger log.Logger - tools_file string - tools_files []string - tools_folder string - prebuiltConfig string - inStream io.Reader - outStream io.Writer - errStream io.Writer + cfg server.ServerConfig + logger log.Logger + tools_file string + tools_files []string + tools_folder string + prebuiltConfigs []string + inStream io.Reader + outStream io.Writer + errStream io.Writer } // NewCommand returns a Command object representing an invocation of the CLI. @@ -376,10 +376,10 @@ func NewCommand(opts ...Option) *Command { flags.StringVar(&cmd.cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.") // Fetch prebuilt tools sources to customize the help description prebuiltHelp := fmt.Sprintf( - "Use a prebuilt tool configuration by source type. Allowed: '%s'.", + "Use a prebuilt tool configuration by source type. Allowed: '%s'. Can be specified multiple times.", strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"), ) - flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", prebuiltHelp) + flags.StringSliceVar(&cmd.prebuiltConfigs, "prebuilt", []string{}, prebuiltHelp) flags.BoolVar(&cmd.cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.") flags.BoolVar(&cmd.cfg.DisableReload, "disable-reload", false, "Disables dynamic reloading of tools file.") flags.BoolVar(&cmd.cfg.UI, "ui", false, "Launches the Toolbox UI web server.") @@ -867,24 +867,32 @@ func run(cmd *Command) error { var allToolsFiles []ToolsFile // Load Prebuilt Configuration - if cmd.prebuiltConfig != "" { - buf, err := prebuiltconfigs.Get(cmd.prebuiltConfig) - if err != nil { - cmd.logger.ErrorContext(ctx, err.Error()) - return err - } - logMsg := fmt.Sprint("Using prebuilt tool configuration for ", cmd.prebuiltConfig) - cmd.logger.InfoContext(ctx, logMsg) - // Append prebuilt.source to Version string for the User Agent - cmd.cfg.Version += "+prebuilt." + cmd.prebuiltConfig - parsed, err := parseToolsFile(ctx, buf) - if err != nil { - errMsg := fmt.Errorf("unable to parse prebuilt tool configuration: %w", err) - cmd.logger.ErrorContext(ctx, errMsg.Error()) - return errMsg + if len(cmd.prebuiltConfigs) > 0 { + slices.Sort(cmd.prebuiltConfigs) + sourcesList := strings.Join(cmd.prebuiltConfigs, ", ") + logMsg := fmt.Sprintf("Using prebuilt tool configurations for: %s", sourcesList) + cmd.logger.InfoContext(ctx, logMsg) + + for _, configName := range cmd.prebuiltConfigs { + buf, err := prebuiltconfigs.Get(configName) + if err != nil { + cmd.logger.ErrorContext(ctx, err.Error()) + return err + } + + // Update version string + cmd.cfg.Version += "+prebuilt." + configName + + // Parse into ToolsFile struct + parsed, err := parseToolsFile(ctx, buf) + if err != nil { + errMsg := fmt.Errorf("unable to parse prebuilt tool configuration for '%s': %w", configName, err) + cmd.logger.ErrorContext(ctx, errMsg.Error()) + return errMsg + } + allToolsFiles = append(allToolsFiles, parsed) } - allToolsFiles = append(allToolsFiles, parsed) } // Determine if Custom Files should be loaded @@ -892,7 +900,7 @@ func run(cmd *Command) error { isCustomConfigured := cmd.tools_file != "" || len(cmd.tools_files) > 0 || cmd.tools_folder != "" // Determine if default 'tools.yaml' should be used (No prebuilt AND No custom flags) - useDefaultToolsFile := cmd.prebuiltConfig == "" && !isCustomConfigured + useDefaultToolsFile := len(cmd.prebuiltConfigs) == 0 && !isCustomConfigured if useDefaultToolsFile { cmd.tools_file = "tools.yaml" diff --git a/cmd/root_test.go b/cmd/root_test.go index e825a36366e..17058d18ff6 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -430,17 +430,27 @@ func TestPrebuiltFlag(t *testing.T) { tcs := []struct { desc string args []string - want string + want []string }{ { desc: "default value", args: []string{}, - want: "", + want: []string{}, }, { - desc: "custom pre built flag", - args: []string{"--tools-file", "alloydb"}, - want: "alloydb", + desc: "single prebuilt flag", + args: []string{"--prebuilt", "alloydb"}, + want: []string{"alloydb"}, + }, + { + desc: "multiple prebuilt flags", + args: []string{"--prebuilt", "alloydb", "--prebuilt", "bigquery"}, + want: []string{"alloydb", "bigquery"}, + }, + { + desc: "comma separated prebuilt flags", + args: []string{"--prebuilt", "alloydb,bigquery"}, + want: []string{"alloydb", "bigquery"}, }, } for _, tc := range tcs { @@ -449,8 +459,8 @@ func TestPrebuiltFlag(t *testing.T) { if err != nil { t.Fatalf("unexpected error invoking command: %s", err) } - if c.tools_file != tc.want { - t.Fatalf("got %v, want %v", c.cfg, tc.want) + if diff := cmp.Diff(c.prebuiltConfigs, tc.want); diff != "" { + t.Fatalf("got %v, want %v, diff %s", c.prebuiltConfigs, tc.want, diff) } }) } @@ -2073,6 +2083,12 @@ authSources: return nil }, }, + { + desc: "sqlite called twice error", + args: []string{"--prebuilt", "sqlite", "--prebuilt", "sqlite"}, + wantErr: true, + errString: "resource conflicts detected", + }, { desc: "tool conflict error", args: []string{"--prebuilt", "sqlite", "--tools-file", toolConflictFile}, diff --git a/docs/en/reference/cli.md b/docs/en/reference/cli.md index a1e6a57985e..686dbc0c739 100644 --- a/docs/en/reference/cli.md +++ b/docs/en/reference/cli.md @@ -16,7 +16,7 @@ description: > | | `--log-level` | Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'. | `info` | | | `--logging-format` | Specify logging format to use. Allowed: 'standard' or 'JSON'. | `standard` | | `-p` | `--port` | Port the server will listen on. | `5000` | -| | `--prebuilt` | Use a prebuilt tool configuration by source type. See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values. | | +| | `--prebuilt` | Use one or more prebuilt tool configuration by source type. See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values. | | | | `--stdio` | Listens via MCP STDIO instead of acting as a remote HTTP server. | | | | `--telemetry-gcp` | Enable exporting directly to Google Cloud Monitoring. | | | | `--telemetry-otlp` | Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318') | | @@ -51,6 +51,11 @@ description: > # Server with prebuilt + custom tools configurations ./toolbox --tools-file tools.yaml --prebuilt alloydb-postgres + +# Server with multiple prebuilt tools configurations +./toolbox --prebuilt alloydb-postgres,alloydb-postgres-admin +# OR +./toolbox --prebuilt alloydb-postgres --prebuilt alloydb-postgres-admin ``` ### Tool Configuration Sources @@ -71,7 +76,7 @@ The CLI supports multiple mutually exclusive ways to specify tool configurations **Prebuilt Configurations:** -- `--prebuilt`: Use predefined configurations for specific database types (e.g., +- `--prebuilt`: Use one or more predefined configurations for specific database types (e.g., 'bigquery', 'postgres', 'spanner'). See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values. diff --git a/docs/en/reference/prebuilt-tools.md b/docs/en/reference/prebuilt-tools.md index d4f347193ba..a539af9fb22 100644 --- a/docs/en/reference/prebuilt-tools.md +++ b/docs/en/reference/prebuilt-tools.md @@ -16,6 +16,9 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP. {{< notice tip >}} You can now use `--prebuilt` along `--tools-file`, `--tools-files`, or `--tools-folder` to combine prebuilt configs with custom tools. + +You can also combine multiple prebuilt configs. + See [Usage Examples](../reference/cli.md#examples). {{< /notice >}}