feat: Support combining multiple prebuilt configurations

This change allows the agent developer to control the maxium number of
rows returned from tools running BigQuery SQL query. Using this feature
the agent developer could limit how large output is presented to LLM in
an agentic user journey.

## Description

> Should include a concise description of the changes (bug or feature),
it's
> impact, along with a summary of the solution

## 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)
- [ ] Make sure to open an issue
https://github.com/googleapis/genai-toolbox/issues/2261
  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)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #2261 2261

DEVELOPER.md

@@ -379,6 +379,23 @@ to approve PRs for main. TeamSync is used to create this team from the MDB
 Group `toolbox-contributors`. Googlers who are developing for MCP-Toolbox
 but aren't part of the core team should join this group.
 
+### Issue/PR Triage and SLO
+After an issue is created, maintainers will assign the following labels:
+* `Priority` (defaulted to P0)
+* `Type` (if applicable)
+* `Product` (if applicable)
+
+All incoming issues and PRs will follow the following SLO:
+| Type            | Priority | Objective                                                              |
+|-----------------|----------|------------------------------------------------------------------------|
+| Feature Request | P0       | Must respond within **5 days**                                         |
+| Process         | P0       | Must respond within **5 days**                                         |
+| Bugs            | P0       | Must respond within **5 days**, and resolve/closure within **14 days** |
+| Bugs            | P1       | Must respond within **7 days**, and resolve/closure within **90 days** |
+| Bugs            | P2       | Must respond within **30 days**
+
+_Types that are not listed in the table do not adhere to any SLO._
+
 ### Releasing
 
 Toolbox has two types of releases: versioned and continuous. It uses Google

cmd/root.go

@@ -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"

cmd/root_test.go

@@ -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},

docs/en/resources/sources/bigquery.md

@@ -134,6 +134,7 @@ sources:
     # scopes: # Optional: List of OAuth scopes to request.
     #   - "https://www.googleapis.com/auth/bigquery"
     #   - "https://www.googleapis.com/auth/drive.readonly"
+    # maxQueryResultRows: 50 # Optional: Limits the number of rows returned by queries. Defaults to 50.
 ```
 
 Initialize a BigQuery source that uses the client's access token:
@@ -153,6 +154,7 @@ sources:
     # scopes: # Optional: List of OAuth scopes to request.
     #   - "https://www.googleapis.com/auth/bigquery"
     #   - "https://www.googleapis.com/auth/drive.readonly"
+    # maxQueryResultRows: 50 # Optional: Limits the number of rows returned by queries. Defaults to 50.
 ```
 
 ## Reference
@@ -167,3 +169,4 @@ sources:
 | useClientOAuth            |   bool   |    false     | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries. **Note:** This cannot be used with `writeMode: protected`.                                                                                                                                                                                                                                                                                                                                                |
 | scopes                    | []string |    false     | A list of OAuth 2.0 scopes to use for the credentials. If not provided, default scopes are used.                                                                                                                                                                                                                                                                                                                                                                                                                     |
 | impersonateServiceAccount |  string  |    false     | Service account email to impersonate when making BigQuery and Dataplex API calls. The authenticated principal must have the `roles/iam.serviceAccountTokenCreator` role on the target service account. [Learn More](https://cloud.google.com/iam/docs/service-account-impersonation)                                                                                                                                                                                                                                |
+| maxQueryResultRows             |   int    |    false     | The maximum number of rows to return from a query. Defaults to 50. |

internal/prebuiltconfigs/tools/bigquery.yaml

@@ -19,6 +19,7 @@ sources:
     location: ${BIGQUERY_LOCATION:}
     useClientOAuth: ${BIGQUERY_USE_CLIENT_OAUTH:false}
     scopes: ${BIGQUERY_SCOPES:}
+    maxQueryResultRows: ${BIGQUERY_MAX_QUERY_RESULT_ROWS:50}
 
 tools:
   analyze_contribution:

internal/sources/bigquery/bigquery.go

@@ -89,6 +89,7 @@ type Config struct {
 	UseClientOAuth            bool                `yaml:"useClientOAuth"`
 	ImpersonateServiceAccount string              `yaml:"impersonateServiceAccount"`
 	Scopes                    StringOrStringSlice `yaml:"scopes"`
+	MaxQueryResultRows        int                 `yaml:"maxQueryResultRows"`
 }
 
 // StringOrStringSlice is a custom type that can unmarshal both a single string
@@ -127,6 +128,10 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 		r.WriteMode = WriteModeAllowed
 	}
 
+	if r.MaxQueryResultRows == 0 {
+		r.MaxQueryResultRows = 50
+	}
+
 	if r.WriteMode == WriteModeProtected && r.UseClientOAuth {
 		// The protected mode only allows write operations to the session's temporary datasets.
 		// when using client OAuth, a new session is created every
@@ -150,7 +155,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 		Client:             client,
 		RestService:        restService,
 		TokenSource:        tokenSource,
-		MaxQueryResultRows: 50,
+		MaxQueryResultRows: r.MaxQueryResultRows,
 		ClientCreator:      clientCreator,
 	}
 
@@ -567,7 +572,7 @@ func (s *Source) RunSQL(ctx context.Context, bqClient *bigqueryapi.Client, state
 	}
 
 	var out []any
-	for {
+	for s.MaxQueryResultRows <= 0 || len(out) < s.MaxQueryResultRows {
 		var val []bigqueryapi.Value
 		err = it.Next(&val)
 		if err == iterator.Done {

internal/sources/bigquery/bigquery_test.go

@@ -21,9 +21,12 @@ import (
 
 	yaml "github.com/goccy/go-yaml"
 	"github.com/google/go-cmp/cmp"
+	"go.opentelemetry.io/otel/trace/noop"
+
 	"github.com/googleapis/genai-toolbox/internal/server"
 	"github.com/googleapis/genai-toolbox/internal/sources/bigquery"
 	"github.com/googleapis/genai-toolbox/internal/testutils"
+	"github.com/googleapis/genai-toolbox/internal/util"
 )
 
 func TestParseFromYamlBigQuery(t *testing.T) {
@@ -154,6 +157,26 @@ func TestParseFromYamlBigQuery(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "with max query result rows example",
+			in: `
+			sources:
+				my-instance:
+					kind: bigquery
+					project: my-project
+					location: us
+					maxQueryResultRows: 10
+			`,
+			want: server.SourceConfigs{
+				"my-instance": bigquery.Config{
+					Name:               "my-instance",
+					Kind:               bigquery.SourceKind,
+					Project:            "my-project",
+					Location:           "us",
+					MaxQueryResultRows: 10,
+				},
+			},
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
@@ -220,6 +243,59 @@ func TestFailParseFromYaml(t *testing.T) {
 	}
 }
 
+func TestInitialize_MaxQueryResultRows(t *testing.T) {
+	ctx, err := testutils.ContextWithNewLogger()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	ctx = util.WithUserAgent(ctx, "test-agent")
+	tracer := noop.NewTracerProvider().Tracer("")
+
+	tcs := []struct {
+		desc string
+		cfg  bigquery.Config
+		want int
+	}{
+		{
+			desc: "default value",
+			cfg: bigquery.Config{
+				Name:           "test-default",
+				Kind:           bigquery.SourceKind,
+				Project:        "test-project",
+				UseClientOAuth: true,
+			},
+			want: 50,
+		},
+		{
+			desc: "configured value",
+			cfg: bigquery.Config{
+				Name:               "test-configured",
+				Kind:               bigquery.SourceKind,
+				Project:            "test-project",
+				UseClientOAuth:     true,
+				MaxQueryResultRows: 100,
+			},
+			want: 100,
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			src, err := tc.cfg.Initialize(ctx, tracer)
+			if err != nil {
+				t.Fatalf("Initialize failed: %v", err)
+			}
+			bqSrc, ok := src.(*bigquery.Source)
+			if !ok {
+				t.Fatalf("Expected *bigquery.Source, got %T", src)
+			}
+			if bqSrc.MaxQueryResultRows != tc.want {
+				t.Errorf("MaxQueryResultRows = %d, want %d", bqSrc.MaxQueryResultRows, tc.want)
+			}
+		})
+	}
+}
+
 func TestNormalizeValue(t *testing.T) {
 	tests := []struct {
 		name     string