[pull] main from googleapis:main

cmd/internal/persistent_flags.go → cmd/internal/flags.go

@@ -20,27 +20,43 @@ import (
 
 	"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 )
 
 // PersistentFlags sets up flags that are available for all commands and
 // subcommands
 // It is also used to set up persistent flags during subcommand unit tests
 func PersistentFlags(parentCmd *cobra.Command, opts *ToolboxOptions) {
 	persistentFlags := parentCmd.PersistentFlags()
-
-	persistentFlags.StringVar(&opts.ToolsFile, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
-	persistentFlags.StringSliceVar(&opts.ToolsFiles, "tools-files", []string{}, "Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with --tools-file, or --tools-folder.")
-	persistentFlags.StringVar(&opts.ToolsFolder, "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(&opts.Cfg.LogLevel, "log-level", "Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'.")
 	persistentFlags.Var(&opts.Cfg.LoggingFormat, "logging-format", "Specify logging format to use. Allowed: 'standard' or 'JSON'.")
 	persistentFlags.BoolVar(&opts.Cfg.TelemetryGCP, "telemetry-gcp", false, "Enable exporting directly to Google Cloud Monitoring.")
 	persistentFlags.StringVar(&opts.Cfg.TelemetryOTLP, "telemetry-otlp", "", "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318')")
 	persistentFlags.StringVar(&opts.Cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.")
+	persistentFlags.StringSliceVar(&opts.Cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
+}
+
+// ConfigFileFlags defines flags related to the configuration file.
+// It should be applied to any command that requires configuration loading.
+func ConfigFileFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
+	flags.StringVar(&opts.ToolsFile, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
+	flags.StringSliceVar(&opts.ToolsFiles, "tools-files", []string{}, "Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with --tools-file, or --tools-folder.")
+	flags.StringVar(&opts.ToolsFolder, "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.")
 	// 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(&opts.PrebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
-	persistentFlags.StringSliceVar(&opts.Cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
+	flags.StringSliceVar(&opts.PrebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
+}
+
+// ServeFlags defines flags for starting and configuring the server.
+func ServeFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
+	flags.StringVarP(&opts.Cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.")
+	flags.IntVarP(&opts.Cfg.Port, "port", "p", 5000, "Port the server will listen on.")
+	flags.BoolVar(&opts.Cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")
+	flags.BoolVar(&opts.Cfg.UI, "ui", false, "Launches the Toolbox UI web server.")
+
+	flags.StringSliceVar(&opts.Cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
+	flags.StringSliceVar(&opts.Cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
 }

cmd/internal/invoke/command.go

@@ -39,6 +39,8 @@ Example:
 			return runInvoke(c, args, opts)
 		},
 	}
+	flags := cmd.Flags()
+	internal.ConfigFileFlags(flags, opts)
 	return cmd
 }
 

cmd/internal/serve/command.go

@@ -0,0 +1,132 @@
+// 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 serve
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/googleapis/genai-toolbox/cmd/internal"
+	"github.com/googleapis/genai-toolbox/internal/server"
+	"github.com/spf13/cobra"
+)
+
+func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "serve",
+		Short: "Deploy the toolbox server",
+		Long:  "Deploy the toolbox server",
+	}
+	flags := cmd.Flags()
+	internal.ServeFlags(flags, opts)
+	cmd.RunE = func(*cobra.Command, []string) error { return runServe(cmd, opts) }
+	return cmd
+}
+
+func runServe(cmd *cobra.Command, opts *internal.ToolboxOptions) 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:
+			opts.Logger.DebugContext(sCtx, "Received SIGINT signal to shutdown.")
+		case syscall.SIGTERM:
+			opts.Logger.DebugContext(sCtx, "Received SIGTERM signal to shutdown.")
+		}
+		cancel()
+	}(ctx)
+
+	ctx, shutdown, err := opts.Setup(ctx)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_ = shutdown(ctx)
+	}()
+
+	// start server
+	s, err := server.NewServer(ctx, opts.Cfg)
+	if err != nil {
+		errMsg := fmt.Errorf("toolbox failed to initialize: %w", err)
+		opts.Logger.ErrorContext(ctx, errMsg.Error())
+		return errMsg
+	}
+
+	// run server in background
+	srvErr := make(chan error, 1)
+	if opts.Cfg.Stdio {
+		go func() {
+			defer close(srvErr)
+			err = s.ServeStdio(ctx, opts.IOStreams.In, opts.IOStreams.Out)
+			if err != nil {
+				srvErr <- err
+			}
+		}()
+	} else {
+		err = s.Listen(ctx)
+		if err != nil {
+			errMsg := fmt.Errorf("toolbox failed to start listener: %w", err)
+			opts.Logger.ErrorContext(ctx, errMsg.Error())
+			return errMsg
+		}
+		opts.Logger.InfoContext(ctx, "Server ready to serve!")
+		if opts.Cfg.UI {
+			opts.Logger.InfoContext(ctx, fmt.Sprintf("Toolbox UI is up and running at: http://%s:%d/ui", opts.Cfg.Address, opts.Cfg.Port))
+		}
+
+		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)
+			opts.Logger.ErrorContext(ctx, errMsg.Error())
+			return errMsg
+		}
+	case <-ctx.Done():
+		shutdownContext, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		opts.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
+}

cmd/internal/serve/command_test.go

@@ -0,0 +1,63 @@
+// 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 serve
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/googleapis/genai-toolbox/cmd/internal"
+	"github.com/spf13/cobra"
+)
+
+func serveCommand(ctx context.Context, args []string) (string, error) {
+	parentCmd := &cobra.Command{Use: "toolbox"}
+
+	buf := new(bytes.Buffer)
+	opts := internal.NewToolboxOptions(internal.WithIOStreams(buf, buf))
+	internal.PersistentFlags(parentCmd, opts)
+
+	cmd := NewCommand(opts)
+	parentCmd.AddCommand(cmd)
+	parentCmd.SetArgs(args)
+	// Inject the context into the Cobra command
+	parentCmd.SetContext(ctx)
+
+	err := parentCmd.Execute()
+	return buf.String(), err
+}
+
+func TestServe(t *testing.T) {
+	// context will automatically shutdown in 1 second.
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	args := []string{"serve", "--port", "0"}
+	output, err := serveCommand(ctx, args)
+	if err != nil {
+		t.Fatalf("expected graceful shutdown without error, got: %v", err)
+	}
+
+	if !strings.Contains(output, "Server ready to serve!") {
+		t.Errorf("expected to find server ready message in output, got: %s", output)
+	}
+
+	if !strings.Contains(output, "Shutting down gracefully...") {
+		t.Errorf("expected to find graceful shutdown message in output, got: %s", output)
+	}
+}

cmd/internal/skills/command.go

@@ -54,13 +54,14 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
 		},
 	}
 
-	cmd.Flags().StringVar(&cmd.name, "name", "", "Name of the generated skill.")
-	cmd.Flags().StringVar(&cmd.description, "description", "", "Description of the generated skill")
-	cmd.Flags().StringVar(&cmd.toolset, "toolset", "", "Name of the toolset to convert into a skill. If not provided, all tools will be included.")
-	cmd.Flags().StringVar(&cmd.outputDir, "output-dir", "skills", "Directory to output generated skills")
-	cmd.Flags().StringVar(&cmd.licenseHeader, "license-header", "", "Optional license header to prepend to generated node scripts.")
-	cmd.Flags().StringVar(&cmd.additionalNotes, "additional-notes", "", "Additional notes to add under the Usage section of the generated SKILL.md")
-
+	flags := cmd.Flags()
+	internal.ConfigFileFlags(flags, opts)
+	flags.StringVar(&cmd.name, "name", "", "Name of the generated skill.")
+	flags.StringVar(&cmd.description, "description", "", "Description of the generated skill")
+	flags.StringVar(&cmd.toolset, "toolset", "", "Name of the toolset to convert into a skill. If not provided, all tools will be included.")
+	flags.StringVar(&cmd.outputDir, "output-dir", "skills", "Directory to output generated skills")
+	flags.StringVar(&cmd.licenseHeader, "license-header", "", "Optional license header to prepend to generated node scripts.")
+	flags.StringVar(&cmd.additionalNotes, "additional-notes", "", "Additional notes to add under the Usage section of the generated SKILL.md")
 	_ = cmd.MarkFlagRequired("name")
 	_ = cmd.MarkFlagRequired("description")
 	return cmd.Command

cmd/root.go

@@ -33,6 +33,7 @@ import (
 	// Importing the cmd/internal package also import packages for side effect of registration
 	"github.com/googleapis/genai-toolbox/cmd/internal"
 	"github.com/googleapis/genai-toolbox/cmd/internal/invoke"
+	"github.com/googleapis/genai-toolbox/cmd/internal/serve"
 	"github.com/googleapis/genai-toolbox/cmd/internal/skills"
 	"github.com/googleapis/genai-toolbox/internal/auth"
 	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
@@ -110,30 +111,21 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
 
 	// setup flags that are common across all commands
 	internal.PersistentFlags(cmd, opts)
-
 	flags := cmd.Flags()
-
-	flags.StringVarP(&opts.Cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.")
-	flags.IntVarP(&opts.Cfg.Port, "port", "p", 5000, "Port the server will listen on.")
-
+	internal.ConfigFileFlags(flags, opts)
+	internal.ServeFlags(flags, opts)
 	flags.StringVar(&opts.ToolsFile, "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")
-	flags.BoolVar(&opts.Cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")
 	flags.BoolVar(&opts.Cfg.DisableReload, "disable-reload", false, "Disables dynamic reloading of tools file.")
-	flags.BoolVar(&opts.Cfg.UI, "ui", false, "Launches the Toolbox UI web server.")
-	// TODO: Insecure by default. Might consider updating this for v1.0.0
-	flags.StringSliceVar(&opts.Cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
-	flags.StringSliceVar(&opts.Cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
 	flags.IntVar(&opts.Cfg.PollInterval, "poll-interval", 0, "Specifies the polling frequency (seconds) for configuration file updates.")
-
 	// wrap RunE command so that we have access to original Command object
 	cmd.RunE = func(*cobra.Command, []string) error { return run(cmd, opts) }
 
-	// Register subcommands for tool invocation
+	// Register subcommands
 	cmd.AddCommand(invoke.NewCommand(opts))
-	// Register subcommands for skill generation
 	cmd.AddCommand(skills.NewCommand(opts))
+	cmd.AddCommand(serve.NewCommand(opts))
 
 	return cmd
 }

cmd/root_test.go

@@ -937,6 +937,7 @@ func TestSubcommandWiring(t *testing.T) {
 	}{
 		{[]string{"invoke"}, "invoke"},
 		{[]string{"skills-generate"}, "skills-generate"},
+		{[]string{"serve"}, "serve"},
 	}
 
 	for _, tc := range tests {

go.mod

@@ -48,6 +48,7 @@ require (
 	github.com/sijms/go-ora/v2 v2.9.0
 	github.com/snowflakedb/gosnowflake v1.19.0
 	github.com/spf13/cobra v1.10.2
+	github.com/spf13/pflag v1.0.9
 	github.com/testcontainers/testcontainers-go v0.40.0
 	github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0
 	github.com/thlib/go-timezone-local v0.0.7
@@ -229,7 +230,6 @@ require (
 	github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
 	github.com/shirou/gopsutil/v4 v4.25.6 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
-	github.com/spf13/pflag v1.0.9 // indirect
 	github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect
 	github.com/stretchr/testify v1.11.1 // indirect
 	github.com/tklauser/go-sysconf v0.3.12 // indirect