Compare commits

..

1 Commits

Author SHA1 Message Date
Disha Prakash
d833d1992d feat: Support combining multiple prebuilt configurations 2026-01-12 21:00:26 +00:00
4 changed files with 59 additions and 41 deletions

View File

@@ -313,15 +313,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.
@@ -374,10 +374,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.")
@@ -864,24 +864,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
@@ -889,7 +897,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"

View File

@@ -420,17 +420,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 {
@@ -439,8 +449,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)
}
})
}
@@ -2063,6 +2073,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},

View File

@@ -101,9 +101,6 @@ authentication token.
TOOLBOX_URL = "https://your-toolbox-service-xyz.a.run.app"
# Initialize the client with the Cloud Run URL and Auth headers
# WARNING: There is a known issue that may cause it to fail with a 401 error
# on local machines.
# See https://github.com/googleapis/mcp-toolbox-sdk-python/issues/496
client = ToolboxSyncClient(
TOOLBOX_URL,
client_headers={"Authorization": auth_methods.get_google_id_token(TOOLBOX_URL)}

View File

@@ -211,9 +211,6 @@ from toolbox_core import ToolboxClient, auth_methods
# Replace with the Cloud Run service URL generated in the previous step
URL = "https://cloud-run-url.app"
# WARNING: There is a known issue that may cause it to fail with a 401 error
# on local machines.
# See https://github.com/googleapis/mcp-toolbox-sdk-python/issues/496
auth_token_provider = auth_methods.aget_google_id_token(URL) # can also use sync method
async def main():