Merge branch 'main' into err

cmd/invoke_tool_test.go

@@ -1,131 +0,0 @@
-// Copyright 2026 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package cmd
-
-import (
-	"context"
-	"os"
-	"path/filepath"
-	"strings"
-	"testing"
-)
-
-func TestInvokeTool(t *testing.T) {
-	// Create a temporary tools file
-	tmpDir := t.TempDir()
-
-	toolsFileContent := `
-sources:
-  my-sqlite:
-    kind: sqlite
-    database: test.db
-tools:
-  hello-sqlite:
-    kind: sqlite-sql
-    source: my-sqlite
-    description: "hello tool"
-    statement: "SELECT 'hello' as greeting"
-  echo-tool:
-    kind: sqlite-sql
-    source: my-sqlite
-    description: "echo tool"
-    statement: "SELECT ? as msg"
-    parameters:
-      - name: message
-        type: string
-        description: message to echo
-`
-
-	toolsFilePath := filepath.Join(tmpDir, "tools.yaml")
-	if err := os.WriteFile(toolsFilePath, []byte(toolsFileContent), 0644); err != nil {
-		t.Fatalf("failed to write tools file: %v", err)
-	}
-
-	tcs := []struct {
-		desc    string
-		args    []string
-		want    string
-		wantErr bool
-		errStr  string
-	}{
-		{
-			desc: "success - basic tool call",
-			args: []string{"invoke", "hello-sqlite", "--tools-file", toolsFilePath},
-			want: `"greeting": "hello"`,
-		},
-		{
-			desc: "success - tool call with parameters",
-			args: []string{"invoke", "echo-tool", `{"message": "world"}`, "--tools-file", toolsFilePath},
-			want: `"msg": "world"`,
-		},
-		{
-			desc:    "error - tool not found",
-			args:    []string{"invoke", "non-existent", "--tools-file", toolsFilePath},
-			wantErr: true,
-			errStr:  `tool "non-existent" not found`,
-		},
-		{
-			desc:    "error - invalid JSON params",
-			args:    []string{"invoke", "echo-tool", `invalid-json`, "--tools-file", toolsFilePath},
-			wantErr: true,
-			errStr:  `params must be a valid JSON string`,
-		},
-	}
-
-	for _, tc := range tcs {
-		t.Run(tc.desc, func(t *testing.T) {
-			_, got, err := invokeCommandWithContext(context.Background(), tc.args)
-			if (err != nil) != tc.wantErr {
-				t.Fatalf("got error %v, wantErr %v", err, tc.wantErr)
-			}
-			if tc.wantErr && !strings.Contains(err.Error(), tc.errStr) {
-				t.Fatalf("got error %v, want error containing %q", err, tc.errStr)
-			}
-			if !tc.wantErr && !strings.Contains(got, tc.want) {
-				t.Fatalf("got %q, want it to contain %q", got, tc.want)
-			}
-		})
-	}
-}
-
-func TestInvokeTool_AuthUnsupported(t *testing.T) {
-	tmpDir := t.TempDir()
-	toolsFileContent := `
-sources:
-  my-bq:
-    kind: bigquery
-    project: my-project
-    useClientOAuth: true
-tools:
-  bq-tool:
-    kind: bigquery-sql
-    source: my-bq
-    description: "bq tool"
-    statement: "SELECT 1"
-`
-	toolsFilePath := filepath.Join(tmpDir, "auth_tools.yaml")
-	if err := os.WriteFile(toolsFilePath, []byte(toolsFileContent), 0644); err != nil {
-		t.Fatalf("failed to write tools file: %v", err)
-	}
-
-	args := []string{"invoke", "bq-tool", "--tools-file", toolsFilePath}
-	_, _, err := invokeCommandWithContext(context.Background(), args)
-	if err == nil {
-		t.Fatal("expected error for tool requiring client auth, but got nil")
-	}
-	if !strings.Contains(err.Error(), "client authorization is not supported") {
-		t.Fatalf("unexpected error message: %v", err)
-	}
-}

cmd/root.go

@@ -34,7 +34,6 @@ import (
 	"github.com/fsnotify/fsnotify"
 	yaml "github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/auth"
-	"github.com/googleapis/genai-toolbox/internal/cli/invoke"
 	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
 	"github.com/googleapis/genai-toolbox/internal/log"
 	"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
@@ -366,42 +365,37 @@ func NewCommand(opts ...Option) *Command {
 	baseCmd.SetErr(cmd.errStream)
 
 	flags := cmd.Flags()
-	persistentFlags := cmd.PersistentFlags()
-
 	flags.StringVarP(&cmd.cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.")
 	flags.IntVarP(&cmd.cfg.Port, "port", "p", 5000, "Port the server will listen on.")
 
 	flags.StringVar(&cmd.tools_file, "tools_file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
 	// deprecate tools_file
 	_ = flags.MarkDeprecated("tools_file", "please use --tools-file instead")
-	persistentFlags.StringVar(&cmd.tools_file, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
-	persistentFlags.StringSliceVar(&cmd.tools_files, "tools-files", []string{}, "Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with --tools-file, or --tools-folder.")
-	persistentFlags.StringVar(&cmd.tools_folder, "tools-folder", "", "Directory path containing YAML tool configuration files. All .yaml and .yml files in the directory will be loaded and merged. Cannot be used with --tools-file, or --tools-files.")
-	persistentFlags.Var(&cmd.cfg.LogLevel, "log-level", "Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'.")
-	persistentFlags.Var(&cmd.cfg.LoggingFormat, "logging-format", "Specify logging format to use. Allowed: 'standard' or 'JSON'.")
-	persistentFlags.BoolVar(&cmd.cfg.TelemetryGCP, "telemetry-gcp", false, "Enable exporting directly to Google Cloud Monitoring.")
-	persistentFlags.StringVar(&cmd.cfg.TelemetryOTLP, "telemetry-otlp", "", "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318')")
-	persistentFlags.StringVar(&cmd.cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.")
+	flags.StringVar(&cmd.tools_file, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
+	flags.StringSliceVar(&cmd.tools_files, "tools-files", []string{}, "Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with --tools-file, or --tools-folder.")
+	flags.StringVar(&cmd.tools_folder, "tools-folder", "", "Directory path containing YAML tool configuration files. All .yaml and .yml files in the directory will be loaded and merged. Cannot be used with --tools-file, or --tools-files.")
+	flags.Var(&cmd.cfg.LogLevel, "log-level", "Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'.")
+	flags.Var(&cmd.cfg.LoggingFormat, "logging-format", "Specify logging format to use. Allowed: 'standard' or 'JSON'.")
+	flags.BoolVar(&cmd.cfg.TelemetryGCP, "telemetry-gcp", false, "Enable exporting directly to Google Cloud Monitoring.")
+	flags.StringVar(&cmd.cfg.TelemetryOTLP, "telemetry-otlp", "", "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318')")
+	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'. Can be specified multiple times.",
 		strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"),
 	)
-	persistentFlags.StringSliceVar(&cmd.prebuiltConfigs, "prebuilt", []string{}, 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.")
 	// TODO: Insecure by default. Might consider updating this for v1.0.0
 	flags.StringSliceVar(&cmd.cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
 	flags.StringSliceVar(&cmd.cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
-	persistentFlags.StringSliceVar(&cmd.cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
+	flags.StringSliceVar(&cmd.cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
 
 	// wrap RunE command so that we have access to original Command object
 	cmd.RunE = func(*cobra.Command, []string) error { return run(cmd) }
 
-	// Register subcommands for tool invocation
-	baseCmd.AddCommand(invoke.NewCommand(cmd))
-
 	return cmd
 }
 
@@ -925,24 +919,71 @@ func resolveWatcherInputs(toolsFile string, toolsFiles []string, toolsFolder str
 	return watchDirs, watchedFiles
 }
 
-func (cmd *Command) Config() server.ServerConfig {
-	return cmd.cfg
-}
+func run(cmd *Command) error {
+	ctx, cancel := context.WithCancel(cmd.Context())
+	defer cancel()
 
-func (cmd *Command) Out() io.Writer {
-	return cmd.outStream
-}
+	// watch for sigterm / sigint signals
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)
+	go func(sCtx context.Context) {
+		var s os.Signal
+		select {
+		case <-sCtx.Done():
+			// this should only happen when the context supplied when testing is canceled
+			return
+		case s = <-signals:
+		}
+		switch s {
+		case syscall.SIGINT:
+			cmd.logger.DebugContext(sCtx, "Received SIGINT signal to shutdown.")
+		case syscall.SIGTERM:
+			cmd.logger.DebugContext(sCtx, "Sending SIGTERM signal to shutdown.")
+		}
+		cancel()
+	}(ctx)
 
-func (cmd *Command) Logger() log.Logger {
-	return cmd.logger
-}
-
-func (cmd *Command) LoadConfig(ctx context.Context) error {
-	logger, err := util.LoggerFromContext(ctx)
-	if err != nil {
-		return err
+	// If stdio, set logger's out stream (usually DEBUG and INFO logs) to errStream
+	loggerOut := cmd.outStream
+	if cmd.cfg.Stdio {
+		loggerOut = cmd.errStream
 	}
 
+	// Handle logger separately from config
+	switch strings.ToLower(cmd.cfg.LoggingFormat.String()) {
+	case "json":
+		logger, err := log.NewStructuredLogger(loggerOut, cmd.errStream, cmd.cfg.LogLevel.String())
+		if err != nil {
+			return fmt.Errorf("unable to initialize logger: %w", err)
+		}
+		cmd.logger = logger
+	case "standard":
+		logger, err := log.NewStdLogger(loggerOut, cmd.errStream, cmd.cfg.LogLevel.String())
+		if err != nil {
+			return fmt.Errorf("unable to initialize logger: %w", err)
+		}
+		cmd.logger = logger
+	default:
+		return fmt.Errorf("logging format invalid")
+	}
+
+	ctx = util.WithLogger(ctx, cmd.logger)
+
+	// Set up OpenTelemetry
+	otelShutdown, err := telemetry.SetupOTel(ctx, cmd.cfg.Version, cmd.cfg.TelemetryOTLP, cmd.cfg.TelemetryGCP, cmd.cfg.TelemetryServiceName)
+	if err != nil {
+		errMsg := fmt.Errorf("error setting up OpenTelemetry: %w", err)
+		cmd.logger.ErrorContext(ctx, errMsg.Error())
+		return errMsg
+	}
+	defer func() {
+		err := otelShutdown(ctx)
+		if err != nil {
+			errMsg := fmt.Errorf("error shutting down OpenTelemetry: %w", err)
+			cmd.logger.ErrorContext(ctx, errMsg.Error())
+		}
+	}()
+
 	var allToolsFiles []ToolsFile
 
 	// Load Prebuilt Configuration
@@ -951,12 +992,12 @@ func (cmd *Command) LoadConfig(ctx context.Context) error {
 		slices.Sort(cmd.prebuiltConfigs)
 		sourcesList := strings.Join(cmd.prebuiltConfigs, ", ")
 		logMsg := fmt.Sprintf("Using prebuilt tool configurations for: %s", sourcesList)
-		logger.InfoContext(ctx, logMsg)
+		cmd.logger.InfoContext(ctx, logMsg)
 
 		for _, configName := range cmd.prebuiltConfigs {
 			buf, err := prebuiltconfigs.Get(configName)
 			if err != nil {
-				logger.ErrorContext(ctx, err.Error())
+				cmd.logger.ErrorContext(ctx, err.Error())
 				return err
 			}
 
@@ -964,7 +1005,7 @@ func (cmd *Command) LoadConfig(ctx context.Context) error {
 			parsed, err := parseToolsFile(ctx, buf)
 			if err != nil {
 				errMsg := fmt.Errorf("unable to parse prebuilt tool configuration for '%s': %w", configName, err)
-				logger.ErrorContext(ctx, errMsg.Error())
+				cmd.logger.ErrorContext(ctx, errMsg.Error())
 				return errMsg
 			}
 			allToolsFiles = append(allToolsFiles, parsed)
@@ -990,7 +1031,7 @@ func (cmd *Command) LoadConfig(ctx context.Context) error {
 			(cmd.tools_file != "" && cmd.tools_folder != "") ||
 			(len(cmd.tools_files) > 0 && cmd.tools_folder != "") {
 			errMsg := fmt.Errorf("--tools-file, --tools-files, and --tools-folder flags cannot be used simultaneously")
-			logger.ErrorContext(ctx, errMsg.Error())
+			cmd.logger.ErrorContext(ctx, errMsg.Error())
 			return errMsg
 		}
 
@@ -999,18 +1040,18 @@ func (cmd *Command) LoadConfig(ctx context.Context) error {
 
 		if len(cmd.tools_files) > 0 {
 			// Use tools-files
-			logger.InfoContext(ctx, fmt.Sprintf("Loading and merging %d tool configuration files", len(cmd.tools_files)))
+			cmd.logger.InfoContext(ctx, fmt.Sprintf("Loading and merging %d tool configuration files", len(cmd.tools_files)))
 			customTools, err = loadAndMergeToolsFiles(ctx, cmd.tools_files)
 		} else if cmd.tools_folder != "" {
 			// Use tools-folder
-			logger.InfoContext(ctx, fmt.Sprintf("Loading and merging all YAML files from directory: %s", cmd.tools_folder))
+			cmd.logger.InfoContext(ctx, fmt.Sprintf("Loading and merging all YAML files from directory: %s", cmd.tools_folder))
 			customTools, err = loadAndMergeToolsFolder(ctx, cmd.tools_folder)
 		} else {
 			// Use single file (tools-file or default `tools.yaml`)
 			buf, readFileErr := os.ReadFile(cmd.tools_file)
 			if readFileErr != nil {
 				errMsg := fmt.Errorf("unable to read tool file at %q: %w", cmd.tools_file, readFileErr)
-				logger.ErrorContext(ctx, errMsg.Error())
+				cmd.logger.ErrorContext(ctx, errMsg.Error())
 				return errMsg
 			}
 			customTools, err = parseToolsFile(ctx, buf)
@@ -1020,7 +1061,7 @@ func (cmd *Command) LoadConfig(ctx context.Context) error {
 		}
 
 		if err != nil {
-			logger.ErrorContext(ctx, err.Error())
+			cmd.logger.ErrorContext(ctx, err.Error())
 			return err
 		}
 		allToolsFiles = append(allToolsFiles, customTools)
@@ -1042,7 +1083,7 @@ func (cmd *Command) LoadConfig(ctx context.Context) error {
 	// This will error if custom tools collide with prebuilt tools
 	finalToolsFile, err := mergeToolsFiles(allToolsFiles...)
 	if err != nil {
-		logger.ErrorContext(ctx, err.Error())
+		cmd.logger.ErrorContext(ctx, err.Error())
 		return err
 	}
 
@@ -1053,91 +1094,15 @@ func (cmd *Command) LoadConfig(ctx context.Context) error {
 	cmd.cfg.ToolsetConfigs = finalToolsFile.Toolsets
 	cmd.cfg.PromptConfigs = finalToolsFile.Prompts
 
-	return nil
-}
-
-func (cmd *Command) Setup(ctx context.Context) (context.Context, func(context.Context) error, error) {
-	// If stdio, set logger's out stream (usually DEBUG and INFO logs) to errStream
-	loggerOut := cmd.outStream
-	if cmd.cfg.Stdio {
-		loggerOut = cmd.errStream
-	}
-
-	// Handle logger separately from config
-	logger, err := log.NewLogger(cmd.cfg.LoggingFormat.String(), cmd.cfg.LogLevel.String(), loggerOut, cmd.errStream)
-	if err != nil {
-		return ctx, nil, fmt.Errorf("unable to initialize logger: %w", err)
-	}
-	cmd.logger = logger
-
-	ctx = util.WithLogger(ctx, cmd.logger)
-
-	// Set up OpenTelemetry
-	otelShutdown, err := telemetry.SetupOTel(ctx, cmd.cfg.Version, cmd.cfg.TelemetryOTLP, cmd.cfg.TelemetryGCP, cmd.cfg.TelemetryServiceName)
-	if err != nil {
-		errMsg := fmt.Errorf("error setting up OpenTelemetry: %w", err)
-		cmd.logger.ErrorContext(ctx, errMsg.Error())
-		return ctx, nil, errMsg
-	}
-
-	shutdownFunc := func(ctx context.Context) error {
-		err := otelShutdown(ctx)
-		if err != nil {
-			errMsg := fmt.Errorf("error shutting down OpenTelemetry: %w", err)
-			cmd.logger.ErrorContext(ctx, errMsg.Error())
-			return err
-		}
-		return nil
-	}
-
-	instrumentation, err := telemetry.CreateTelemetryInstrumentation(cmd.cfg.Version)
+	instrumentation, err := telemetry.CreateTelemetryInstrumentation(versionString)
 	if err != nil {
 		errMsg := fmt.Errorf("unable to create telemetry instrumentation: %w", err)
 		cmd.logger.ErrorContext(ctx, errMsg.Error())
-		return ctx, shutdownFunc, errMsg
+		return errMsg
 	}
 
 	ctx = util.WithInstrumentation(ctx, instrumentation)
 
-	return ctx, shutdownFunc, nil
-}
-
-func run(cmd *Command) error {
-	ctx, cancel := context.WithCancel(cmd.Context())
-	defer cancel()
-
-	// watch for sigterm / sigint signals
-	signals := make(chan os.Signal, 1)
-	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)
-	go func(sCtx context.Context) {
-		var s os.Signal
-		select {
-		case <-sCtx.Done():
-			// this should only happen when the context supplied when testing is canceled
-			return
-		case s = <-signals:
-		}
-		switch s {
-		case syscall.SIGINT:
-			cmd.logger.DebugContext(sCtx, "Received SIGINT signal to shutdown.")
-		case syscall.SIGTERM:
-			cmd.logger.DebugContext(sCtx, "Sending SIGTERM signal to shutdown.")
-		}
-		cancel()
-	}(ctx)
-
-	ctx, shutdown, err := cmd.Setup(ctx)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		_ = shutdown(ctx)
-	}()
-
-	if err := cmd.LoadConfig(ctx); err != nil {
-		return err
-	}
-
 	// start server
 	s, err := server.NewServer(ctx, cmd.cfg)
 	if err != nil {
@@ -1177,9 +1142,6 @@ func run(cmd *Command) error {
 		}()
 	}
 
-	// Determine if Custom Files are configured (re-check as loadAndMergeConfig might have updated defaults)
-	isCustomConfigured := cmd.tools_file != "" || len(cmd.tools_files) > 0 || cmd.tools_folder != ""
-
 	if isCustomConfigured && !cmd.cfg.DisableReload {
 		watchDirs, watchedFiles := resolveWatcherInputs(cmd.tools_file, cmd.tools_files, cmd.tools_folder)
 		// start watching the file(s) or folder for changes to trigger dynamic reloading

docs/ALLOYDBADMIN_README.md

@@ -23,14 +23,6 @@ To connect to the database to explore and query data, search the MCP store for t
 
 In the Antigravity MCP Store, click the "Install" button.
 
-> [!NOTE]
-> On first use, the installation process automatically downloads and uses
-> [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-> `>=0.26.0`. To update MCP Toolbox, use:
-> ```npm i -g @toolbox-sdk/server@latest```
-> To always run the latest version, update the MCP server configuration to use:
-> ```npx -y @toolbox-sdk/server@latest --prebuilt alloydb-postgres-admin```.
-
 You'll now be able to see all enabled tools in the "Tools" tab.
 
 > [!NOTE]

docs/ALLOYDBPG_README.md

@@ -27,13 +27,6 @@ For AlloyDB infrastructure management, search the MCP store for the AlloyDB for
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt alloydb-postgres```.
 
 2. Add the required inputs for your [cluster](https://docs.cloud.google.com/alloydb/docs/cluster-list) in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/BIGQUERY_README.md

@@ -21,13 +21,6 @@ An editor configured to use the BigQuery MCP server can use its AI capabilities
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt bigquery```.
 
 2. Add the required inputs in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/CLOUDSQLMSSQLADMIN_README.md

@@ -23,14 +23,6 @@ To connect to the database to explore and query data, search the MCP store for t
 
 In the Antigravity MCP Store, click the "Install" button.
 
-> [!NOTE]
-> On first use, the installation process automatically downloads and uses
-> [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-> `>=0.26.0`. To update MCP Toolbox, use:
-> ```npm i -g @toolbox-sdk/server@latest```
-> To always run the latest version, update the MCP server configuration to use:
-> ```npx -y @toolbox-sdk/server@latest --prebuilt cloud-sql-mssql-admin```.
-
 You'll now be able to see all enabled tools in the "Tools" tab.
 
 > [!NOTE]

docs/CLOUDSQLMSSQL_README.md

@@ -24,13 +24,6 @@ For Cloud SQL infrastructure management, search the MCP store for the Cloud SQL
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt cloud-sql-mssql```.
 
 2. Add the required inputs for your [instance](https://cloud.google.com/sql/docs/sqlserver/instance-info) in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/CLOUDSQLMYSQLADMIN_README.md

@@ -23,14 +23,6 @@ To connect to the database to explore and query data, search the MCP store for t
 
 In the Antigravity MCP Store, click the "Install" button.
 
-> [!NOTE]
-> On first use, the installation process automatically downloads and uses
-> [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-> `>=0.26.0`. To update MCP Toolbox, use:
-> ```npm i -g @toolbox-sdk/server@latest```
-> To always run the latest version, update the MCP server configuration to use:
-> ```npx -y @toolbox-sdk/server@latest --prebuilt cloud-sql-mysql-admin```.
-
 You'll now be able to see all enabled tools in the "Tools" tab.
 
 > [!NOTE]

docs/CLOUDSQLMYSQL_README.md

@@ -26,13 +26,6 @@ For Cloud SQL infrastructure management, search the MCP store for the Cloud SQL
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt cloud-sql-mysql```.
 
 2. Add the required inputs for your [instance](https://cloud.google.com/sql/docs/mysql/instance-info) in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/CLOUDSQLPGADMIN_README.md

@@ -23,14 +23,6 @@ To connect to the database to explore and query data, search the MCP store for t
 
 In the Antigravity MCP Store, click the "Install" button.
 
-> [!NOTE]
-> On first use, the installation process automatically downloads and uses
-> [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-> `>=0.26.0`. To update MCP Toolbox, use:
-> ```npm i -g @toolbox-sdk/server@latest```
-> To always run the latest version, update the MCP server configuration to use:
-> ```npx -y @toolbox-sdk/server@latest --prebuilt cloud-sql-postgres-admin```.
-
 You'll now be able to see all enabled tools in the "Tools" tab.
 
 > [!NOTE]

docs/CLOUDSQLPG_README.md

@@ -26,13 +26,6 @@ For Cloud SQL infrastructure management, search the MCP store for the Cloud SQL
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt cloud-sql-postgres```.
 
 2. Add the required inputs for your [instance](https://cloud.google.com/sql/docs/postgres/instance-info) in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/DATAPLEX_README.md

@@ -20,13 +20,6 @@ An editor configured to use the Dataplex MCP server can use its AI capabilities
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt dataplex```.
 
 2. Add the required inputs in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/LOOKER_README.md

@@ -21,13 +21,6 @@ An editor configured to use the Looker MCP server can use its AI capabilities to
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt looker```.
 
 2. Add the required inputs for your [instance](https://docs.cloud.google.com/looker/docs/set-up-and-administer-looker) in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/SPANNER_README.md

@@ -21,13 +21,6 @@ An editor configured to use the Cloud Spanner MCP server can use its AI capabili
 ## Install & Configuration
 
 1. In the Antigravity MCP Store, click the "Install" button.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest --prebuilt spanner```.
 
 2. Add the required inputs for your [instance](https://docs.cloud.google.com/spanner/docs/instances) in the configuration pop-up, then click "Save". You can update this configuration at any time in the "Configure" tab.
 

docs/TOOLBOX_README.md

@@ -12,17 +12,10 @@ The MCP Toolbox for Databases Server gives AI-powered development tools the abil
 ## Install & Configuration
 
 1.  In the Antigravity MCP Store, click the **Install** button. A configuration window will appear.
-    > [!NOTE]
-    > On first use, the installation process automatically downloads and uses
-    > [MCP Toolbox](https://www.npmjs.com/package/@toolbox-sdk/server)
-    > `>=0.26.0`. To update MCP Toolbox, use:
-    > ```npm i -g @toolbox-sdk/server@latest```
-    > To always run the latest version, update the MCP server configuration to use:
-    > ```npx -y @toolbox-sdk/server@latest```. 
 
-3.  Create your [`tools.yaml` configuration file](https://googleapis.github.io/genai-toolbox/getting-started/configure/).
+2.  Create your [`tools.yaml` configuration file](https://googleapis.github.io/genai-toolbox/getting-started/configure/).
 
-4.  In the configuration window, enter the full absolute path to your `tools.yaml` file and click **Save**.
+3.  In the configuration window, enter the full absolute path to your `tools.yaml` file and click **Save**.
 
 > [!NOTE]
 > If you encounter issues with Windows Defender blocking the execution, you may need to configure an allowlist. See [Configure exclusions for Microsoft Defender Antivirus](https://learn.microsoft.com/en-us/microsoft-365/security/defender-endpoint/configure-exclusions-microsoft-defender-antivirus?view=o365-worldwide) for more details.

docs/en/how-to/invoke_tool.md

@@ -1,74 +0,0 @@
----
-title: "Invoke Tools via CLI"
-type: docs
-weight: 10
-description: >
-  Learn how to invoke your tools directly from the command line using the `invoke` command.
----
-
-The `invoke` command allows you to invoke tools defined in your configuration directly from the CLI. This is useful for:
-
-- **Ephemeral Invocation:** Executing a tool without spinning up a full MCP server/client.
-- **Debugging:** Isolating tool execution logic and testing with various parameter combinations.
-
-{{< notice tip >}}
-**Keep configurations minimal:** The `invoke` command initializes *all* resources (sources, tools, etc.) defined in your configuration files during execution. To ensure fast response times, consider using a minimal configuration file containing only the tools you need for the specific invocation.
-{{< notice tip >}}
-
-## Prerequisites
-
-- You have the `toolbox` binary installed or built.
-- You have a valid tool configuration file (e.g., `tools.yaml`).
-
-## Basic Usage
-
-The basic syntax for the command is:
-
-```bash
-toolbox [--tools-file <path> | --prebuilt <name>] invoke <tool-name> [params]
-```
-
-- `<tool-name>`: The name of the tool you want to call. This must match the name defined in your `tools.yaml`.
-- `[params]`: (Optional) A JSON string representing the arguments for the tool.
-
-## Examples
-
-### 1. Calling a Tool without Parameters
-
-If your tool takes no parameters, simply provide the tool name:
-
-```bash
-toolbox  --tools-file tools.yaml invoke my-simple-tool
-```
-
-### 2. Calling a Tool with Parameters
-
-For tools that require arguments, pass them as a JSON string. Ensure you escape quotes correctly for your shell.
-
-**Example: A tool that takes parameters**
-
-Assuming a tool named `mytool` taking `a` and `b`:
-
-```bash
-toolbox --tools-file tools.yaml invoke mytool '{"a": 10, "b": 20}' 
-```
-
-**Example: A tool that queries a database**
-
-```bash
-toolbox  --tools-file tools.yaml invoke db-query '{"sql": "SELECT * FROM users LIMIT 5"}'
-```
-
-### 3. Using Prebuilt Configurations
-
-You can also use the `--prebuilt` flag to load prebuilt toolsets.
-
-```bash
-toolbox --prebuilt cloudsql-postgres invoke cloudsql-postgres-list-instances
-```
-
-## Troubleshooting
-
-- **Tool not found:** Ensure the `<tool-name>` matches exactly what is in your YAML file and that the file is correctly loaded via `--tools-file`.
-- **Invalid parameters:** Double-check your JSON syntax. The error message will usually indicate if the JSON parsing failed or if the parameters didn't match the tool's schema.
-- **Auth errors:** The `invoke` command currently does not support flows requiring client-side authorization (like OAuth flow initiation via the CLI). It works best for tools using service-side authentication (e.g., Application Default Credentials).

docs/en/reference/cli.md

@@ -30,21 +30,6 @@ description: >
 |              | `--user-agent-metadata`    | Appends additional metadata to the User-Agent.                                                                                                                                   |             |
 | `-v`         | `--version`                | version for toolbox                                                                                                                                                              |             |
 
-## Sub Commands
-
-### `invoke`
-
-Executes a tool directly with the provided parameters. This is useful for testing tool configurations and parameters without needing a full client setup.
-
-**Syntax:**
-
-```bash
-toolbox invoke <tool-name> [params]
-```
-
-- `<tool-name>`: The name of the tool to execute (as defined in your configuration).
-- `[params]`: (Optional) A JSON string containing the parameters for the tool.
-
 ## Examples
 
 ### Transport Configuration

go.mod

@@ -52,7 +52,7 @@ require (
 	github.com/trinodb/trino-go-client v0.330.0
 	github.com/valkey-io/valkey-go v1.0.68
 	github.com/yugabyte/pgx/v5 v5.5.3-yb-5
-	go.mongodb.org/mongo-driver/v2 v2.4.2
+	go.mongodb.org/mongo-driver v1.17.4
 	go.opentelemetry.io/contrib/propagators/autoprop v0.62.0
 	go.opentelemetry.io/otel v1.38.0
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v1.38.0
@@ -211,6 +211,7 @@ require (
 	github.com/moby/term v0.5.2 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/montanaflynn/stats v0.7.1 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
 	github.com/mtibben/percent v0.2.1 // indirect
 	github.com/nakagami/chacha20 v0.1.0 // indirect

go.sum

@@ -1238,6 +1238,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe/go.mod h1:wL8QJuTMNUDYhXwkmfOly8iTdp5TEcJFWZD2D7SIkUc=
+github.com/montanaflynn/stats v0.7.1 h1:etflOAAHORrCC44V+aR6Ftzort912ZU+YLiSTuV8eaE=
+github.com/montanaflynn/stats v0.7.1/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
 github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
 github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
 github.com/mtibben/percent v0.2.1 h1:5gssi8Nqo8QU/r2pynCm+hBQHpkB/uNK7BJCFogWdzs=
@@ -1402,8 +1404,8 @@ github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaD
 gitlab.com/nyarla/go-crypt v0.0.0-20160106005555-d9a5dc2b789b h1:7gd+rd8P3bqcn/96gOZa3F5dpJr/vEiDQYlNb/y2uNs=
 gitlab.com/nyarla/go-crypt v0.0.0-20160106005555-d9a5dc2b789b/go.mod h1:T3BPAOm2cqquPa0MKWeNkmOM5RQsRhkrwMWonFMN7fE=
 go.mongodb.org/mongo-driver v1.11.4/go.mod h1:PTSz5yu21bkT/wXpkS7WR5f0ddqw5quethTUn9WM+2g=
-go.mongodb.org/mongo-driver/v2 v2.4.2 h1:HrJ+Auygxceby9MLp3YITobef5a8Bv4HcPFIkml1U7U=
-go.mongodb.org/mongo-driver/v2 v2.4.2/go.mod h1:jHeEDJHJq7tm6ZF45Issun9dbogjfnPySb1vXA7EeAI=
+go.mongodb.org/mongo-driver v1.17.4 h1:jUorfmVzljjr0FLzYQsGP8cgN/qzzxlY9Vh0C9KFXVw=
+go.mongodb.org/mongo-driver v1.17.4/go.mod h1:Hy04i7O2kC4RS06ZrhPRqj/u4DTYkFDAAccj+rVKqgQ=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8=
 go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw=

internal/cli/invoke/command.go

@@ -1,161 +0,0 @@
-// Copyright 2026 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package invoke
-
-import (
-	"context"
-	"encoding/json"
-	"fmt"
-	"io"
-
-	"github.com/googleapis/genai-toolbox/internal/log"
-	"github.com/googleapis/genai-toolbox/internal/server"
-	"github.com/googleapis/genai-toolbox/internal/server/resources"
-	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"github.com/spf13/cobra"
-)
-
-// RootCommand defines the interface for required by invoke subcommand.
-// This allows subcommands to access shared resources and functionality without
-// direct coupling to the root command's implementation.
-type RootCommand interface {
-	// Config returns a copy of the current server configuration.
-	Config() server.ServerConfig
-
-	// Out returns the writer used for standard output.
-	Out() io.Writer
-
-	// LoadConfig loads and merges the configuration from files, folders, and prebuilts.
-	LoadConfig(ctx context.Context) error
-
-	// Setup initializes the runtime environment, including logging and telemetry.
-	// It returns the updated context and a shutdown function to be called when finished.
-	Setup(ctx context.Context) (context.Context, func(context.Context) error, error)
-
-	// Logger returns the logger instance.
-	Logger() log.Logger
-}
-
-func NewCommand(rootCmd RootCommand) *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "invoke <tool-name> [params]",
-		Short: "Execute a tool directly",
-		Long: `Execute a tool directly with parameters.
-Params must be a JSON string.
-Example:
-  toolbox invoke my-tool '{"param1": "value1"}'`,
-		Args: cobra.MinimumNArgs(1),
-		RunE: func(c *cobra.Command, args []string) error {
-			return runInvoke(c, args, rootCmd)
-		},
-	}
-	return cmd
-}
-
-func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
-	ctx, cancel := context.WithCancel(cmd.Context())
-	defer cancel()
-
-	ctx, shutdown, err := rootCmd.Setup(ctx)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		_ = shutdown(ctx)
-	}()
-
-	// Load and merge tool configurations
-	if err := rootCmd.LoadConfig(ctx); err != nil {
-		return err
-	}
-
-	// Initialize Resources
-	sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, rootCmd.Config())
-	if err != nil {
-		errMsg := fmt.Errorf("failed to initialize resources: %w", err)
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-
-	resourceMgr := resources.NewResourceManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
-
-	// Execute Tool
-	toolName := args[0]
-	tool, ok := resourceMgr.GetTool(toolName)
-	if !ok {
-		errMsg := fmt.Errorf("tool %q not found", toolName)
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-
-	var paramsInput string
-	if len(args) > 1 {
-		paramsInput = args[1]
-	}
-
-	params := make(map[string]any)
-	if paramsInput != "" {
-		if err := json.Unmarshal([]byte(paramsInput), &params); err != nil {
-			errMsg := fmt.Errorf("params must be a valid JSON string: %w", err)
-			rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-			return errMsg
-		}
-	}
-
-	parsedParams, err := parameters.ParseParams(tool.GetParameters(), params, nil)
-	if err != nil {
-		errMsg := fmt.Errorf("invalid parameters: %w", err)
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-
-	parsedParams, err = tool.EmbedParams(ctx, parsedParams, resourceMgr.GetEmbeddingModelMap())
-	if err != nil {
-		errMsg := fmt.Errorf("error embedding parameters: %w", err)
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-
-	// Client Auth not supported for ephemeral CLI call
-	requiresAuth, err := tool.RequiresClientAuthorization(resourceMgr)
-	if err != nil {
-		errMsg := fmt.Errorf("failed to check auth requirements: %w", err)
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-	if requiresAuth {
-		errMsg := fmt.Errorf("client authorization is not supported")
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-
-	result, err := tool.Invoke(ctx, resourceMgr, parsedParams, "")
-	if err != nil {
-		errMsg := fmt.Errorf("tool execution failed: %w", err)
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-
-	// Print Result
-	output, err := json.MarshalIndent(result, "", "  ")
-	if err != nil {
-		errMsg := fmt.Errorf("failed to marshal result: %w", err)
-		rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
-		return errMsg
-	}
-	fmt.Fprintln(rootCmd.Out(), string(output))
-
-	return nil
-}

internal/log/log.go

@@ -22,18 +22,6 @@ import (
 	"strings"
 )
 
-// NewLogger creates a new logger based on the provided format and level.
-func NewLogger(format, level string, out, err io.Writer) (Logger, error) {
-	switch strings.ToLower(format) {
-	case "json":
-		return NewStructuredLogger(out, err, level)
-	case "standard":
-		return NewStdLogger(out, err, level)
-	default:
-		return nil, fmt.Errorf("logging format invalid: %s", format)
-	}
-}
-
 // StdLogger is the standard logger
 type StdLogger struct {
 	outLogger *slog.Logger

internal/sources/mongodb/mongodb.go

@@ -23,9 +23,9 @@ import (
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/util"
-	"go.mongodb.org/mongo-driver/v2/bson"
-	"go.mongodb.org/mongo-driver/v2/mongo"
-	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -148,7 +148,7 @@ func (s *Source) Aggregate(ctx context.Context, pipelineString string, canonical
 	return res, err
 }
 
-func (s *Source) Find(ctx context.Context, filterString, database, collection string, opts *options.FindOptionsBuilder) ([]any, error) {
+func (s *Source) Find(ctx context.Context, filterString, database, collection string, opts *options.FindOptions) ([]any, error) {
 	var filter = bson.D{}
 	err := bson.UnmarshalExtJSON([]byte(filterString), false, &filter)
 	if err != nil {
@@ -163,7 +163,7 @@ func (s *Source) Find(ctx context.Context, filterString, database, collection st
 	return parseData(ctx, cur)
 }
 
-func (s *Source) FindOne(ctx context.Context, filterString, database, collection string, opts *options.FindOneOptionsBuilder) ([]any, error) {
+func (s *Source) FindOne(ctx context.Context, filterString, database, collection string, opts *options.FindOneOptions) ([]any, error) {
 	var filter = bson.D{}
 	err := bson.UnmarshalExtJSON([]byte(filterString), false, &filter)
 	if err != nil {
@@ -233,7 +233,7 @@ func (s *Source) UpdateMany(ctx context.Context, filterString string, canonical
 		return nil, fmt.Errorf("unable to unmarshal update string: %w", err)
 	}
 
-	res, err := s.MongoClient().Database(database).Collection(collection).UpdateMany(ctx, filter, update, options.UpdateMany().SetUpsert(upsert))
+	res, err := s.MongoClient().Database(database).Collection(collection).UpdateMany(ctx, filter, update, options.Update().SetUpsert(upsert))
 	if err != nil {
 		return nil, fmt.Errorf("error updating collection: %w", err)
 	}
@@ -252,7 +252,7 @@ func (s *Source) UpdateOne(ctx context.Context, filterString string, canonical b
 		return nil, fmt.Errorf("unable to unmarshal update string: %w", err)
 	}
 
-	res, err := s.MongoClient().Database(database).Collection(collection).UpdateOne(ctx, filter, update, options.UpdateOne().SetUpsert(upsert))
+	res, err := s.MongoClient().Database(database).Collection(collection).UpdateOne(ctx, filter, update, options.Update().SetUpsert(upsert))
 	if err != nil {
 		return nil, fmt.Errorf("error updating collection: %w", err)
 	}
@@ -266,7 +266,7 @@ func (s *Source) DeleteMany(ctx context.Context, filterString, database, collect
 		return nil, err
 	}
 
-	res, err := s.MongoClient().Database(database).Collection(collection).DeleteMany(ctx, filter, options.DeleteMany())
+	res, err := s.MongoClient().Database(database).Collection(collection).DeleteMany(ctx, filter, options.Delete())
 	if err != nil {
 		return nil, err
 	}
@@ -284,7 +284,7 @@ func (s *Source) DeleteOne(ctx context.Context, filterString, database, collecti
 		return nil, err
 	}
 
-	res, err := s.MongoClient().Database(database).Collection(collection).DeleteOne(ctx, filter, options.DeleteOne())
+	res, err := s.MongoClient().Database(database).Collection(collection).DeleteOne(ctx, filter, options.Delete())
 	if err != nil {
 		return nil, err
 	}
@@ -303,7 +303,7 @@ func initMongoDBClient(ctx context.Context, tracer trace.Tracer, name, uri strin
 
 	// Create a new MongoDB client
 	clientOpts := options.Client().ApplyURI(uri).SetAppName(userAgent)
-	client, err := mongo.Connect(clientOpts)
+	client, err := mongo.Connect(ctx, clientOpts)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create MongoDB client: %w", err)
 	}

internal/tools/mongodb/mongodbaggregate/mongodbaggregate.go

@@ -21,7 +21,7 @@ import (
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/mongo"
 
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"

internal/tools/mongodb/mongodbdeletemany/mongodbdeletemany.go

@@ -21,7 +21,7 @@ import (
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/mongo"
 
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"

internal/tools/mongodb/mongodbdeleteone/mongodbdeleteone.go

@@ -21,7 +21,7 @@ import (
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/mongo"
 
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"

internal/tools/mongodb/mongodbfind/mongodbfind.go

@@ -22,9 +22,9 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
 	"github.com/googleapis/genai-toolbox/internal/util"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/bson"
-	"go.mongodb.org/mongo-driver/v2/mongo"
-	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
 
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"
@@ -48,7 +48,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
 
 type compatibleSource interface {
 	MongoClient() *mongo.Client
-	Find(context.Context, string, string, string, *options.FindOptionsBuilder) ([]any, error)
+	Find(context.Context, string, string, string, *options.FindOptions) ([]any, error)
 }
 
 type Config struct {
@@ -118,7 +118,7 @@ type Tool struct {
 	mcpManifest tools.McpManifest
 }
 
-func getOptions(ctx context.Context, sortParameters parameters.Parameters, projectPayload string, limit int64, paramsMap map[string]any) (*options.FindOptionsBuilder, error) {
+func getOptions(ctx context.Context, sortParameters parameters.Parameters, projectPayload string, limit int64, paramsMap map[string]any) (*options.FindOptions, error) {
 	logger, err := util.LoggerFromContext(ctx)
 	if err != nil {
 		panic(err)

internal/tools/mongodb/mongodbfindone/mongodbfindone.go

@@ -21,9 +21,9 @@ import (
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/bson"
-	"go.mongodb.org/mongo-driver/v2/mongo"
-	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
 
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"
@@ -47,7 +47,7 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
 
 type compatibleSource interface {
 	MongoClient() *mongo.Client
-	FindOne(context.Context, string, string, string, *options.FindOneOptionsBuilder) ([]any, error)
+	FindOne(context.Context, string, string, string, *options.FindOneOptions) ([]any, error)
 }
 
 type Config struct {

internal/tools/mongodb/mongodbinsertmany/mongodbinsertmany.go

@@ -23,7 +23,7 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 const resourceType string = "mongodb-insert-many"

internal/tools/mongodb/mongodbinsertone/mongodbinsertone.go

@@ -23,7 +23,7 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 const resourceType string = "mongodb-insert-one"

internal/tools/mongodb/mongodbupdatemany/mongodbupdatemany.go

@@ -23,7 +23,7 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 const resourceType string = "mongodb-update-many"

internal/tools/mongodb/mongodbupdateone/mongodbupdateone.go

@@ -23,7 +23,7 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/sources"
 	"github.com/googleapis/genai-toolbox/internal/tools"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 const resourceType string = "mongodb-update-one"

internal/util/errors.go

@@ -0,0 +1,60 @@
+// Copyright 2026 Google LLC
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package util
+
+import "fmt"
+
+type ErrorCategory string
+
+const (
+	CategoryAgent  ErrorCategory = "AGENT_ERROR"
+	CategoryServer ErrorCategory = "SERVER_ERROR"
+)
+
+// ToolboxError is the interface all custom errors must satisfy
+type ToolboxError interface {
+	error
+	Category() ErrorCategory
+}
+
+// Agent Errors
+type AgentError struct {
+	Msg   string
+	Cause error
+}
+
+func (e *AgentError) Error() string { return e.Msg }
+
+func (e *AgentError) Category() ErrorCategory { return CategoryAgent }
+
+func (e *AgentError) Unwrap() error { return e.Cause }
+
+func NewAgentError(msg string, args ...any) *AgentError {
+	return &AgentError{Msg: fmt.Sprintf(msg, args...)}
+}
+
+// Server Errors
+type ServerError struct {
+	Msg   string
+	Cause error
+}
+
+func (e *ServerError) Error() string { return fmt.Sprintf("%s: %v", e.Msg, e.Cause) }
+
+func (e *ServerError) Category() ErrorCategory { return CategoryServer }
+
+func (e *ServerError) Unwrap() error { return e.Cause }
+
+func NewServerError(msg string, cause error) *ServerError {
+	return &ServerError{Msg: msg, Cause: cause}
+}

tests/alloydbpg/alloydb_pg_integration_test.go

@@ -139,25 +139,12 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
 
 	// set up data for param tool
 	createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)
-	teardownTable1,err := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
-	if err != nil {
-    // If an error happened, we still try to run teardown if it was returned
-    	if teardownTable1 != nil {
-        	defer teardownTable1(t)
-    	}
-    	t.Fatalf("Setup failed: %v", err)
-	}
+	teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
+	defer teardownTable1(t)
 
 	// set up data for auth tool
 	createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
-	teardownTable2,err := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
-	if err != nil {
-    // If an error happened, we still try to run teardown if it was returned
-    	if teardownTable2 != nil {
-        	defer teardownTable2(t)
-    	}
-    	t.Fatalf("Setup failed: %v", err)
-	}	
+	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
 	defer teardownTable2(t)
 
 	// Set up table for semanti search

tests/cloudsqlpg/cloud_sql_pg_integration_test.go

@@ -124,26 +124,12 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
 
 	// set up data for param tool
 	createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)
-	teardownTable1,err := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
-	if err != nil {
-    // If an error happened, we still try to run teardown if it was returned
-    	if teardownTable1 != nil {
-        	defer teardownTable1(t)
-    	}
-    	t.Fatalf("Setup failed: %v", err)
-	}
+	teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
 	defer teardownTable1(t)
 
 	// set up data for auth tool
 	createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
-	teardownTable2,err := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
-	if err != nil {
-    // If an error happened, we still try to run teardown if it was returned
-    	if teardownTable2 != nil {
-        	defer teardownTable2(t)
-    	}
-    	t.Fatalf("Setup failed: %v", err)
-	}	
+	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
 	defer teardownTable2(t)
 
 	// Set up table for semantic search

tests/common.go

@@ -613,37 +613,31 @@ func GetMySQLWants() (string, string, string, string) {
 
 // SetupPostgresSQLTable creates and inserts data into a table of tool
 // compatible with postgres-sql tool
-func SetupPostgresSQLTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool, createStatement, insertStatement, tableName string, params []any) (func(*testing.T), error) {
+func SetupPostgresSQLTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool, createStatement, insertStatement, tableName string, params []any) func(*testing.T) {
 	err := pool.Ping(ctx)
 	if err != nil {
-		// Return nil for the function and the error itself
-		return nil, fmt.Errorf("unable to connect to test database: %w", err)
+		t.Fatalf("unable to connect to test database: %s", err)
 	}
 
 	// Create table
 	_, err = pool.Query(ctx, createStatement)
 	if err != nil {
-		return nil, fmt.Errorf("unable to create test table %s: %w", tableName, err)
+		t.Fatalf("unable to create test table %s: %s", tableName, err)
 	}
 
 	// Insert test data
 	_, err = pool.Query(ctx, insertStatement, params...)
 	if err != nil {
-		// If creation worked but insert failed, you might still want to return 
-		// the teardown so the empty table can be cleaned up
-		teardown := func(t *testing.T) {
-			pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName))
-		}
-		return teardown, fmt.Errorf("unable to insert test data: %w", err)
+		t.Fatalf("unable to insert test data: %s", err)
 	}
 
-	// Return the cleanup function and nil for error
 	return func(t *testing.T) {
-		_, err = pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName))
+		// tear down test
+		_, err = pool.Exec(ctx, fmt.Sprintf("DROP TABLE %s;", tableName))
 		if err != nil {
 			t.Errorf("Teardown failed: %s", err)
 		}
-	}, nil
+	}
 }
 
 // SetupMsSQLTable creates and inserts data into a table of tool

tests/mongodb/mongodb_integration_test.go

@@ -28,8 +28,8 @@ import (
 
 	"github.com/googleapis/genai-toolbox/internal/testutils"
 	"github.com/googleapis/genai-toolbox/tests"
-	"go.mongodb.org/mongo-driver/v2/mongo"
-	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
 )
 
 var (
@@ -55,7 +55,7 @@ func getMongoDBVars(t *testing.T) map[string]any {
 
 func initMongoDbDatabase(ctx context.Context, uri, database string) (*mongo.Database, error) {
 	// Create a new mongodb Database
-	client, err := mongo.Connect(options.Client().ApplyURI(uri))
+	client, err := mongo.Connect(ctx, options.Client().ApplyURI(uri))
 	if err != nil {
 		return nil, fmt.Errorf("unable to connect to mongodb: %s", err)
 	}

tests/postgres/postgres_integration_test.go

@@ -103,27 +103,12 @@ func TestPostgres(t *testing.T) {
 
 	// set up data for param tool
 	createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)
-	// teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
-	teardownTable1, err := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
-	if err != nil {
-    // If an error happened, we still try to run teardown if it was returned
-    	if teardownTable1 != nil {
-        	defer teardownTable1(t)
-    	}
-    	t.Fatalf("Setup failed: %v", err)
-	}
+	teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
 	defer teardownTable1(t)
 
 	// set up data for auth tool
 	createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
-	teardownTable2,err := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
-	if err != nil {
-    // If an error happened, we still try to run teardown if it was returned
-    	if teardownTable2 != nil {
-        	defer teardownTable2(t)
-    	}
-    	t.Fatalf("Setup failed: %v", err)
-	}
+	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
 	defer teardownTable2(t)
 
 	// Set up table for semantic search