genai-toolbox/cmd/root.go

// Copyright 2024 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"
	_ "embed"
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	yaml "github.com/goccy/go-yaml"
	"github.com/googleapis/genai-toolbox/internal/log"
	"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
	"github.com/googleapis/genai-toolbox/internal/server"
	"github.com/googleapis/genai-toolbox/internal/telemetry"
	"github.com/googleapis/genai-toolbox/internal/util"

	// Import tool packages for side effect of registration
	_ "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl"
	_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigqueryexecutesql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerygetdatasetinfo"
	_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerygettableinfo"
	_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerylistdatasetids"
	_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerylisttableids"
	_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerysql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/bigtable"
	_ "github.com/googleapis/genai-toolbox/internal/tools/couchbase"
	_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
	_ "github.com/googleapis/genai-toolbox/internal/tools/http"
	_ "github.com/googleapis/genai-toolbox/internal/tools/mssql/mssqlexecutesql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/mssql/mssqlsql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlexecutesql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlsql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/neo4j"
	_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresexecutesql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgressql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/redis"
	_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerexecutesql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/sqlitesql"
	_ "github.com/googleapis/genai-toolbox/internal/tools/valkey"

	"github.com/spf13/cobra"

	_ "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
	_ "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
	_ "github.com/googleapis/genai-toolbox/internal/sources/bigtable"
	_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql"
	_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
	_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
	_ "github.com/googleapis/genai-toolbox/internal/sources/couchbase"
	_ "github.com/googleapis/genai-toolbox/internal/sources/dgraph"
	_ "github.com/googleapis/genai-toolbox/internal/sources/http"
	_ "github.com/googleapis/genai-toolbox/internal/sources/mssql"
	_ "github.com/googleapis/genai-toolbox/internal/sources/mysql"
	_ "github.com/googleapis/genai-toolbox/internal/sources/neo4j"
	_ "github.com/googleapis/genai-toolbox/internal/sources/postgres"
	_ "github.com/googleapis/genai-toolbox/internal/sources/redis"
	_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
	_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
	_ "github.com/googleapis/genai-toolbox/internal/sources/valkey"
)

var (
	// versionString stores the full semantic version, including build metadata.
	versionString string
	// versionNum indicates the numerical part fo the version
	//go:embed version.txt
	versionNum string
	// metadataString indicates additional build or distribution metadata.
	buildType string = "dev" // should be one of "dev", "binary", or "container"
	// commitSha is the git commit it was built from
	commitSha string
)

func init() {
	versionString = semanticVersion()
}

// semanticVersion returns the version of the CLI including a compile-time metadata.
func semanticVersion() string {
	metadataStrings := []string{buildType, runtime.GOOS, runtime.GOARCH}
	if commitSha != "" {
		metadataStrings = append(metadataStrings, commitSha)
	}
	v := strings.TrimSpace(versionNum) + "+" + strings.Join(metadataStrings, ".")
	return v
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := NewCommand().Execute(); err != nil {
		exit := 1
		os.Exit(exit)
	}
}

// Command represents an invocation of the CLI.
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
}

// NewCommand returns a Command object representing an invocation of the CLI.
func NewCommand(opts ...Option) *Command {
	in := os.Stdin
	out := os.Stdout
	err := os.Stderr

	baseCmd := &cobra.Command{
		Use:           "toolbox",
		Version:       versionString,
		SilenceErrors: true,
	}
	cmd := &Command{
		Command:   baseCmd,
		inStream:  in,
		outStream: out,
		errStream: err,
	}

	for _, o := range opts {
		o(cmd)
	}

	// Set server version
	cmd.cfg.Version = versionString

	// set baseCmd in, out and err the same as cmd.
	baseCmd.SetIn(cmd.inStream)
	baseCmd.SetOut(cmd.outStream)
	baseCmd.SetErr(cmd.errStream)

	flags := cmd.Flags()
	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 --prebuilt.")
	// deprecate tools_file
	_ = flags.MarkDeprecated("tools_file", "please use --tools-file instead")
	flags.StringVar(&cmd.tools_file, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --prebuilt, --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 --prebuilt, --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 --prebuilt, --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.")
	flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", "Use a prebuilt tool configuration by source type. Cannot be used with --tools-file. Allowed: 'alloydb-postgres', 'bigquery', 'cloud-sql-mysql', 'cloud-sql-postgres', 'cloud-sql-mssql', 'postgres', 'spanner', 'spanner-postgres'.")
	flags.BoolVar(&cmd.cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")

	// wrap RunE command so that we have access to original Command object
	cmd.RunE = func(*cobra.Command, []string) error { return run(cmd) }

	return cmd
}

type ToolsFile struct {
	Sources      server.SourceConfigs      `yaml:"sources"`
	AuthSources  server.AuthServiceConfigs `yaml:"authSources"` // Deprecated: Kept for compatibility.
	AuthServices server.AuthServiceConfigs `yaml:"authServices"`
	Tools        server.ToolConfigs        `yaml:"tools"`
	Toolsets     server.ToolsetConfigs     `yaml:"toolsets"`
}

// parseEnv replaces environment variables ${ENV_NAME} with their values.
func parseEnv(input string) string {
	re := regexp.MustCompile(`\$\{(\w+)\}`)

	return re.ReplaceAllStringFunc(input, func(match string) string {
		parts := re.FindStringSubmatch(match)
		if len(parts) < 2 {
			// technically shouldn't happen
			return match
		}

		// extract the variable name
		variableName := parts[1]
		if value, found := os.LookupEnv(variableName); found {
			return value
		}
		return match
	})
}

// parseToolsFile parses the provided yaml into appropriate configs.
func parseToolsFile(ctx context.Context, raw []byte) (ToolsFile, error) {
	var toolsFile ToolsFile
	// Replace environment variables if found
	raw = []byte(parseEnv(string(raw)))
	// Parse contents
	err := yaml.UnmarshalContext(ctx, raw, &toolsFile, yaml.Strict())
	if err != nil {
		return toolsFile, err
	}
	return toolsFile, nil
}

// mergeToolsFiles merges multiple ToolsFile structs into one.
// Detects and raises errors for resource conflicts in sources, authServices, tools, and toolsets.
// All resource names (sources, authServices, tools, toolsets) must be unique across all files.
func mergeToolsFiles(files ...ToolsFile) (ToolsFile, error) {
	merged := ToolsFile{
		Sources:      make(server.SourceConfigs),
		AuthServices: make(server.AuthServiceConfigs),
		Tools:        make(server.ToolConfigs),
		Toolsets:     make(server.ToolsetConfigs),
	}

	var conflicts []string

	for fileIndex, file := range files {
		// Check for conflicts and merge sources
		for name, source := range file.Sources {
			if _, exists := merged.Sources[name]; exists {
				conflicts = append(conflicts, fmt.Sprintf("source '%s' (file #%d)", name, fileIndex+1))
			} else {
				merged.Sources[name] = source
			}
		}

		// Check for conflicts and merge authSources (deprecated, but still support)
		for name, authSource := range file.AuthSources {
			if _, exists := merged.AuthSources[name]; exists {
				conflicts = append(conflicts, fmt.Sprintf("authSource '%s' (file #%d)", name, fileIndex+1))
			} else {
				merged.AuthSources[name] = authSource
			}
		}

		// Check for conflicts and merge authServices
		for name, authService := range file.AuthServices {
			if _, exists := merged.AuthServices[name]; exists {
				conflicts = append(conflicts, fmt.Sprintf("authService '%s' (file #%d)", name, fileIndex+1))
			} else {
				merged.AuthServices[name] = authService
			}
		}

		// Check for conflicts and merge tools
		for name, tool := range file.Tools {
			if _, exists := merged.Tools[name]; exists {
				conflicts = append(conflicts, fmt.Sprintf("tool '%s' (file #%d)", name, fileIndex+1))
			} else {
				merged.Tools[name] = tool
			}
		}

		// Check for conflicts and merge toolsets
		for name, toolset := range file.Toolsets {
			if _, exists := merged.Toolsets[name]; exists {
				conflicts = append(conflicts, fmt.Sprintf("toolset '%s' (file #%d)", name, fileIndex+1))
			} else {
				merged.Toolsets[name] = toolset
			}
		}
	}

	// If conflicts were detected, return an error
	if len(conflicts) > 0 {
		return ToolsFile{}, fmt.Errorf("resource conflicts detected:\n  - %s\n\nPlease ensure each source, authService, tool, and toolset has a unique name across all files", strings.Join(conflicts, "\n  - "))
	}

	return merged, nil
}

// loadAndMergeToolsFiles loads multiple YAML files and merges them
func loadAndMergeToolsFiles(ctx context.Context, filePaths []string) (ToolsFile, error) {
	var toolsFiles []ToolsFile

	for _, filePath := range filePaths {
		buf, err := os.ReadFile(filePath)
		if err != nil {
			return ToolsFile{}, fmt.Errorf("unable to read tool file at %q: %w", filePath, err)
		}

		toolsFile, err := parseToolsFile(ctx, buf)
		if err != nil {
			return ToolsFile{}, fmt.Errorf("unable to parse tool file at %q: %w", filePath, err)
		}

		toolsFiles = append(toolsFiles, toolsFile)
	}

	mergedFile, err := mergeToolsFiles(toolsFiles...)
	if err != nil {
		return ToolsFile{}, fmt.Errorf("unable to merge tools files: %w", err)
	}

	return mergedFile, nil
}

// loadAndMergeToolsFolder loads all YAML files from a directory and merges them
func loadAndMergeToolsFolder(ctx context.Context, folderPath string) (ToolsFile, error) {
	// Check if directory exists
	info, err := os.Stat(folderPath)
	if err != nil {
		return ToolsFile{}, fmt.Errorf("unable to access tools folder at %q: %w", folderPath, err)
	}
	if !info.IsDir() {
		return ToolsFile{}, fmt.Errorf("path %q is not a directory", folderPath)
	}

	// Find all YAML files in the directory
	pattern := filepath.Join(folderPath, "*.yaml")
	yamlFiles, err := filepath.Glob(pattern)
	if err != nil {
		return ToolsFile{}, fmt.Errorf("error finding YAML files in %q: %w", folderPath, err)
	}

	// Also find .yml files
	ymlPattern := filepath.Join(folderPath, "*.yml")
	ymlFiles, err := filepath.Glob(ymlPattern)
	if err != nil {
		return ToolsFile{}, fmt.Errorf("error finding YML files in %q: %w", folderPath, err)
	}

	// Combine both file lists
	allFiles := append(yamlFiles, ymlFiles...)
	
	if len(allFiles) == 0 {
		return ToolsFile{}, fmt.Errorf("no YAML files found in directory %q", folderPath)
	}

	// Use existing loadAndMergeToolsFiles function
	return loadAndMergeToolsFiles(ctx, allFiles)
}

// updateLogLevel checks if Toolbox have to update the existing log level set by users.
// stdio doesn't support "debug" and "info" logs.
func updateLogLevel(stdio bool, logLevel string) bool {
	if stdio {
		switch strings.ToUpper(logLevel) {
		case log.Debug, log.Info:
			return true
		default:
			return false
		}
	}
	return false
}

func run(cmd *Command) error {
	if updateLogLevel(cmd.cfg.Stdio, cmd.cfg.LogLevel.String()) {
		cmd.cfg.LogLevel = server.StringLevel(log.Warn)
	}

	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)

	// Handle logger separately from config
	switch strings.ToLower(cmd.cfg.LoggingFormat.String()) {
	case "json":
		logger, err := log.NewStructuredLogger(cmd.outStream, 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(cmd.outStream, 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 toolsFile ToolsFile

	if cmd.prebuiltConfig != "" {
		// Make sure --prebuilt and --tools-file/--tools-files/--tools-folder flags are mutually exclusive
		if cmd.tools_file != "" || len(cmd.tools_files) > 0 || cmd.tools_folder != "" {
			errMsg := fmt.Errorf("--prebuilt and --tools-file/--tools-files/--tools-folder flags cannot be used simultaneously")
			cmd.logger.ErrorContext(ctx, errMsg.Error())
			return errMsg
		}
		// Use prebuilt tools
		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

		toolsFile, 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
		}
	} else if len(cmd.tools_files) > 0 {
		// Make sure --tools-file, --tools-files, and --tools-folder flags are mutually exclusive
		if cmd.tools_file != "" || cmd.tools_folder != "" {
			errMsg := fmt.Errorf("--tools-file, --tools-files, and --tools-folder flags cannot be used simultaneously")
			cmd.logger.ErrorContext(ctx, errMsg.Error())
			return errMsg
		}
		// Use multiple tools files
		cmd.logger.InfoContext(ctx, fmt.Sprintf("Loading and merging %d tool configuration files", len(cmd.tools_files)))
		var err error
		toolsFile, err = loadAndMergeToolsFiles(ctx, cmd.tools_files)
		if err != nil {
			cmd.logger.ErrorContext(ctx, err.Error())
			return err
		}
	} else if cmd.tools_folder != "" {
		// Make sure --tools-folder and other flags are mutually exclusive
		if cmd.tools_file != "" || len(cmd.tools_files) > 0 {
			errMsg := fmt.Errorf("--tools-file, --tools-files, and --tools-folder flags cannot be used simultaneously")
			cmd.logger.ErrorContext(ctx, errMsg.Error())
			return errMsg
		}
		// Use tools folder
		cmd.logger.InfoContext(ctx, fmt.Sprintf("Loading and merging all YAML files from directory: %s", cmd.tools_folder))
		var err error
		toolsFile, err = loadAndMergeToolsFolder(ctx, cmd.tools_folder)
		if err != nil {
			cmd.logger.ErrorContext(ctx, err.Error())
			return err
		}
	} else {
		// Set default value of tools-file flag to tools.yaml
		if cmd.tools_file == "" {
			cmd.tools_file = "tools.yaml"
		}
		// Read single tool file contents
		buf, err := os.ReadFile(cmd.tools_file)
		if err != nil {
			errMsg := fmt.Errorf("unable to read tool file at %q: %w", cmd.tools_file, err)
			cmd.logger.ErrorContext(ctx, errMsg.Error())
			return errMsg
		}

		toolsFile, err = parseToolsFile(ctx, buf)
		if err != nil {
			errMsg := fmt.Errorf("unable to parse tool file at %q: %w", cmd.tools_file, err)
			cmd.logger.ErrorContext(ctx, errMsg.Error())
			return errMsg
		}
	}

	cmd.cfg.SourceConfigs, cmd.cfg.AuthServiceConfigs, cmd.cfg.ToolConfigs, cmd.cfg.ToolsetConfigs = toolsFile.Sources, toolsFile.AuthServices, toolsFile.Tools, toolsFile.Toolsets
	authSourceConfigs := toolsFile.AuthSources
	if authSourceConfigs != nil {
		cmd.logger.WarnContext(ctx, "`authSources` is deprecated, use `authServices` instead")
		cmd.cfg.AuthServiceConfigs = authSourceConfigs
	}

	// start server
	s, err := server.NewServer(ctx, cmd.cfg, cmd.logger)
	if err != nil {
		errMsg := fmt.Errorf("toolbox failed to initialize: %w", err)
		cmd.logger.ErrorContext(ctx, errMsg.Error())
		return errMsg
	}

	// run server in background
	srvErr := make(chan error)
	if cmd.cfg.Stdio {
		go func() {
			defer close(srvErr)
			err = s.ServeStdio(ctx, cmd.inStream, cmd.outStream)
			if err != nil {
				srvErr <- err
			}
		}()
	} else {
		err = s.Listen(ctx)
		if err != nil {
			errMsg := fmt.Errorf("toolbox failed to start listener: %w", err)
			cmd.logger.ErrorContext(ctx, errMsg.Error())
			return errMsg
		}
		cmd.logger.InfoContext(ctx, "Server ready to serve!")

		go func() {
			defer close(srvErr)
			err = s.Serve(ctx)
			if err != nil {
				srvErr <- err
			}
		}()
	}

	// wait for either the server to error out or the command's context to be canceled
	select {
	case err := <-srvErr:
		if err != nil {
			errMsg := fmt.Errorf("toolbox crashed with the following error: %w", err)
			cmd.logger.ErrorContext(ctx, errMsg.Error())
			return errMsg
		}
	case <-ctx.Done():
		shutdownContext, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		cmd.logger.WarnContext(shutdownContext, "Shutting down gracefully...")
		err := s.Shutdown(shutdownContext)
		if err == context.DeadlineExceeded {
			return fmt.Errorf("graceful shutdown timed out... forcing exit")
		}
	}

	return nil
}