From 1e2c7c7804c67bebf5e2ee9b67c6feb6f05292fd Mon Sep 17 00:00:00 2001 From: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> Date: Tue, 17 Mar 2026 14:46:45 -0700 Subject: [PATCH] feat(cli): add serve subcommand (#2550) The `serve` subcommand **starts a toolbox server** without any primitives. It WILL NOT run with config file. In the future, users could use `serve` with a backend storage. To stop or shutdown the server, user can just terminate the port. The terminate signal will shutdown the server. This new addition **WILL NOT** be a breaking change to existing users. Users can still run toolbox as is. **CLI command:** ``` toolbox serve ``` **Flags associated with the serve subcommand:** | flag | description | default value | | --- | --- | --- | | address | Address of the interface the server will listen on. | 127.0.0.1 | | port | Port the server will listen on. | 5000 | | stdio | Listens via MCP STDIO instead of acting as a remote HTTP server. | false | | ui | Launches the Toolbox UI web server. | false | | allowed-origins | Specifies a list of origins permitted to access this server. | `*` | | allowed-hosts | Specifies a list of hosts permitted to access this server. | `*` | **This PR does the following:** * Add a new `serve` subcommand. Including unit tests for the subcommand * Rename the `cmd/internal/persistent_flags.go` to `cmd/internal/flags.go`, and refactored some flag definitions into dedicated functions. This change allows us to scope flags to specific subcommands as needed, rather than forcing all subcommands to inherit them globally via `PersistentFlags`. --------- Co-authored-by: Averi Kitsch --- .../{persistent_flags.go => flags.go} | 28 +++- cmd/internal/invoke/command.go | 2 + cmd/internal/serve/command.go | 132 ++++++++++++++++++ cmd/internal/serve/command_test.go | 63 +++++++++ cmd/internal/skills/command.go | 15 +- cmd/root.go | 18 +-- cmd/root_test.go | 1 + go.mod | 2 +- 8 files changed, 234 insertions(+), 27 deletions(-) rename cmd/internal/{persistent_flags.go => flags.go} (55%) create mode 100644 cmd/internal/serve/command.go create mode 100644 cmd/internal/serve/command_test.go diff --git a/cmd/internal/persistent_flags.go b/cmd/internal/flags.go similarity index 55% rename from cmd/internal/persistent_flags.go rename to cmd/internal/flags.go index 3874521a15d1..9f2c5100a257 100644 --- a/cmd/internal/persistent_flags.go +++ b/cmd/internal/flags.go @@ -20,6 +20,7 @@ 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 @@ -27,20 +28,35 @@ import ( // 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 '*'.") } diff --git a/cmd/internal/invoke/command.go b/cmd/internal/invoke/command.go index 0dc8cb4d7403..849db710dc09 100644 --- a/cmd/internal/invoke/command.go +++ b/cmd/internal/invoke/command.go @@ -39,6 +39,8 @@ Example: return runInvoke(c, args, opts) }, } + flags := cmd.Flags() + internal.ConfigFileFlags(flags, opts) return cmd } diff --git a/cmd/internal/serve/command.go b/cmd/internal/serve/command.go new file mode 100644 index 000000000000..f57934f07686 --- /dev/null +++ b/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 +} diff --git a/cmd/internal/serve/command_test.go b/cmd/internal/serve/command_test.go new file mode 100644 index 000000000000..06d3e8544e77 --- /dev/null +++ b/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) + } +} diff --git a/cmd/internal/skills/command.go b/cmd/internal/skills/command.go index 4e7f28dfe851..da19d1ddb98b 100644 --- a/cmd/internal/skills/command.go +++ b/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 diff --git a/cmd/root.go b/cmd/root.go index d5c282894a5b..0301d29c2a0c 100644 --- a/cmd/root.go +++ b/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 } diff --git a/cmd/root_test.go b/cmd/root_test.go index 288274a1c543..2ee65837caa9 100644 --- a/cmd/root_test.go +++ b/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 { diff --git a/go.mod b/go.mod index 28f269a9273f..4604935fb206 100644 --- a/go.mod +++ b/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