Skip to content

Commit 1e2c7c7

Browse files
Yuan325averikitsch
Yuan325
and
averikitsch
authored
feat(cli): add serve subcommand (googleapis#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 <akitsch@google.com>
1 parent 1528d7c commit 1e2c7c7

File tree

8 files changed

+234
-27
lines changed

8 files changed

+234
-27
lines changed

cmd/internal/persistent_flags.go renamed to cmd/internal/flags.go

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,27 +20,43 @@ import (
2020

2121
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
2222
"github.com/spf13/cobra"
23+
"github.com/spf13/pflag"
2324
)
2425

2526
// PersistentFlags sets up flags that are available for all commands and
2627
// subcommands
2728
// It is also used to set up persistent flags during subcommand unit tests
2829
func PersistentFlags(parentCmd *cobra.Command, opts *ToolboxOptions) {
2930
persistentFlags := parentCmd.PersistentFlags()
30-
31-
persistentFlags.StringVar(&opts.ToolsFile, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
32-
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.")
33-
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.")
3431
persistentFlags.Var(&opts.Cfg.LogLevel, "log-level", "Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'.")
3532
persistentFlags.Var(&opts.Cfg.LoggingFormat, "logging-format", "Specify logging format to use. Allowed: 'standard' or 'JSON'.")
3633
persistentFlags.BoolVar(&opts.Cfg.TelemetryGCP, "telemetry-gcp", false, "Enable exporting directly to Google Cloud Monitoring.")
3734
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')")
3835
persistentFlags.StringVar(&opts.Cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.")
36+
persistentFlags.StringSliceVar(&opts.Cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
37+
}
38+
39+
// ConfigFileFlags defines flags related to the configuration file.
40+
// It should be applied to any command that requires configuration loading.
41+
func ConfigFileFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
42+
flags.StringVar(&opts.ToolsFile, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
43+
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.")
44+
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.")
3945
// Fetch prebuilt tools sources to customize the help description
4046
prebuiltHelp := fmt.Sprintf(
4147
"Use a prebuilt tool configuration by source type. Allowed: '%s'. Can be specified multiple times.",
4248
strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"),
4349
)
44-
persistentFlags.StringSliceVar(&opts.PrebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
45-
persistentFlags.StringSliceVar(&opts.Cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
50+
flags.StringSliceVar(&opts.PrebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
51+
}
52+
53+
// ServeFlags defines flags for starting and configuring the server.
54+
func ServeFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
55+
flags.StringVarP(&opts.Cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.")
56+
flags.IntVarP(&opts.Cfg.Port, "port", "p", 5000, "Port the server will listen on.")
57+
flags.BoolVar(&opts.Cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")
58+
flags.BoolVar(&opts.Cfg.UI, "ui", false, "Launches the Toolbox UI web server.")
59+
60+
flags.StringSliceVar(&opts.Cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
61+
flags.StringSliceVar(&opts.Cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
4662
}

‎cmd/internal/invoke/command.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ Example:
3939
return runInvoke(c, args, opts)
4040
},
4141
}
42+
flags := cmd.Flags()
43+
internal.ConfigFileFlags(flags, opts)
4244
return cmd
4345
}
4446

‎cmd/internal/serve/command.go‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package serve
16+
17+
import (
18+
"context"
19+
"fmt"
20+
"os"
21+
"os/signal"
22+
"syscall"
23+
"time"
24+
25+
"github.com/googleapis/genai-toolbox/cmd/internal"
26+
"github.com/googleapis/genai-toolbox/internal/server"
27+
"github.com/spf13/cobra"
28+
)
29+
30+
func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
31+
cmd := &cobra.Command{
32+
Use: "serve",
33+
Short: "Deploy the toolbox server",
34+
Long: "Deploy the toolbox server",
35+
}
36+
flags := cmd.Flags()
37+
internal.ServeFlags(flags, opts)
38+
cmd.RunE = func(*cobra.Command, []string) error { return runServe(cmd, opts) }
39+
return cmd
40+
}
41+
42+
func runServe(cmd *cobra.Command, opts *internal.ToolboxOptions) error {
43+
ctx, cancel := context.WithCancel(cmd.Context())
44+
defer cancel()
45+
46+
// watch for sigterm / sigint signals
47+
signals := make(chan os.Signal, 1)
48+
signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)
49+
go func(sCtx context.Context) {
50+
var s os.Signal
51+
select {
52+
case <-sCtx.Done():
53+
// this should only happen when the context supplied when testing is canceled
54+
return
55+
case s = <-signals:
56+
}
57+
switch s {
58+
case syscall.SIGINT:
59+
opts.Logger.DebugContext(sCtx, "Received SIGINT signal to shutdown.")
60+
case syscall.SIGTERM:
61+
opts.Logger.DebugContext(sCtx, "Received SIGTERM signal to shutdown.")
62+
}
63+
cancel()
64+
}(ctx)
65+
66+
ctx, shutdown, err := opts.Setup(ctx)
67+
if err != nil {
68+
return err
69+
}
70+
defer func() {
71+
_ = shutdown(ctx)
72+
}()
73+
74+
// start server
75+
s, err := server.NewServer(ctx, opts.Cfg)
76+
if err != nil {
77+
errMsg := fmt.Errorf("toolbox failed to initialize: %w", err)
78+
opts.Logger.ErrorContext(ctx, errMsg.Error())
79+
return errMsg
80+
}
81+
82+
// run server in background
83+
srvErr := make(chan error, 1)
84+
if opts.Cfg.Stdio {
85+
go func() {
86+
defer close(srvErr)
87+
err = s.ServeStdio(ctx, opts.IOStreams.In, opts.IOStreams.Out)
88+
if err != nil {
89+
srvErr <- err
90+
}
91+
}()
92+
} else {
93+
err = s.Listen(ctx)
94+
if err != nil {
95+
errMsg := fmt.Errorf("toolbox failed to start listener: %w", err)
96+
opts.Logger.ErrorContext(ctx, errMsg.Error())
97+
return errMsg
98+
}
99+
opts.Logger.InfoContext(ctx, "Server ready to serve!")
100+
if opts.Cfg.UI {
101+
opts.Logger.InfoContext(ctx, fmt.Sprintf("Toolbox UI is up and running at: http://%s:%d/ui", opts.Cfg.Address, opts.Cfg.Port))
102+
}
103+
104+
go func() {
105+
defer close(srvErr)
106+
err = s.Serve(ctx)
107+
if err != nil {
108+
srvErr <- err
109+
}
110+
}()
111+
}
112+
113+
// wait for either the server to error out or the command's context to be canceled
114+
select {
115+
case err := <-srvErr:
116+
if err != nil {
117+
errMsg := fmt.Errorf("toolbox crashed with the following error: %w", err)
118+
opts.Logger.ErrorContext(ctx, errMsg.Error())
119+
return errMsg
120+
}
121+
case <-ctx.Done():
122+
shutdownContext, cancel := context.WithTimeout(context.Background(), 10*time.Second)
123+
defer cancel()
124+
opts.Logger.WarnContext(shutdownContext, "Shutting down gracefully...")
125+
err := s.Shutdown(shutdownContext)
126+
if err == context.DeadlineExceeded {
127+
return fmt.Errorf("graceful shutdown timed out... forcing exit")
128+
}
129+
}
130+
131+
return nil
132+
}

‎cmd/internal/serve/command_test.go‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package serve
16+
17+
import (
18+
"bytes"
19+
"context"
20+
"strings"
21+
"testing"
22+
"time"
23+
24+
"github.com/googleapis/genai-toolbox/cmd/internal"
25+
"github.com/spf13/cobra"
26+
)
27+
28+
func serveCommand(ctx context.Context, args []string) (string, error) {
29+
parentCmd := &cobra.Command{Use: "toolbox"}
30+
31+
buf := new(bytes.Buffer)
32+
opts := internal.NewToolboxOptions(internal.WithIOStreams(buf, buf))
33+
internal.PersistentFlags(parentCmd, opts)
34+
35+
cmd := NewCommand(opts)
36+
parentCmd.AddCommand(cmd)
37+
parentCmd.SetArgs(args)
38+
// Inject the context into the Cobra command
39+
parentCmd.SetContext(ctx)
40+
41+
err := parentCmd.Execute()
42+
return buf.String(), err
43+
}
44+
45+
func TestServe(t *testing.T) {
46+
// context will automatically shutdown in 1 second.
47+
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
48+
defer cancel()
49+
50+
args := []string{"serve", "--port", "0"}
51+
output, err := serveCommand(ctx, args)
52+
if err != nil {
53+
t.Fatalf("expected graceful shutdown without error, got: %v", err)
54+
}
55+
56+
if !strings.Contains(output, "Server ready to serve!") {
57+
t.Errorf("expected to find server ready message in output, got: %s", output)
58+
}
59+
60+
if !strings.Contains(output, "Shutting down gracefully...") {
61+
t.Errorf("expected to find graceful shutdown message in output, got: %s", output)
62+
}
63+
}

‎cmd/internal/skills/command.go‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,14 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
5454
},
5555
}
5656

57-
cmd.Flags().StringVar(&cmd.name, "name", "", "Name of the generated skill.")
58-
cmd.Flags().StringVar(&cmd.description, "description", "", "Description of the generated skill")
59-
cmd.Flags().StringVar(&cmd.toolset, "toolset", "", "Name of the toolset to convert into a skill. If not provided, all tools will be included.")
60-
cmd.Flags().StringVar(&cmd.outputDir, "output-dir", "skills", "Directory to output generated skills")
61-
cmd.Flags().StringVar(&cmd.licenseHeader, "license-header", "", "Optional license header to prepend to generated node scripts.")
62-
cmd.Flags().StringVar(&cmd.additionalNotes, "additional-notes", "", "Additional notes to add under the Usage section of the generated SKILL.md")
63-
57+
flags := cmd.Flags()
58+
internal.ConfigFileFlags(flags, opts)
59+
flags.StringVar(&cmd.name, "name", "", "Name of the generated skill.")
60+
flags.StringVar(&cmd.description, "description", "", "Description of the generated skill")
61+
flags.StringVar(&cmd.toolset, "toolset", "", "Name of the toolset to convert into a skill. If not provided, all tools will be included.")
62+
flags.StringVar(&cmd.outputDir, "output-dir", "skills", "Directory to output generated skills")
63+
flags.StringVar(&cmd.licenseHeader, "license-header", "", "Optional license header to prepend to generated node scripts.")
64+
flags.StringVar(&cmd.additionalNotes, "additional-notes", "", "Additional notes to add under the Usage section of the generated SKILL.md")
6465
_ = cmd.MarkFlagRequired("name")
6566
_ = cmd.MarkFlagRequired("description")
6667
return cmd.Command

‎cmd/root.go‎

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
// Importing the cmd/internal package also import packages for side effect of registration
3434
"github.com/googleapis/genai-toolbox/cmd/internal"
3535
"github.com/googleapis/genai-toolbox/cmd/internal/invoke"
36+
"github.com/googleapis/genai-toolbox/cmd/internal/serve"
3637
"github.com/googleapis/genai-toolbox/cmd/internal/skills"
3738
"github.com/googleapis/genai-toolbox/internal/auth"
3839
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
@@ -110,30 +111,21 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
110111

111112
// setup flags that are common across all commands
112113
internal.PersistentFlags(cmd, opts)
113-
114114
flags := cmd.Flags()
115-
116-
flags.StringVarP(&opts.Cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.")
117-
flags.IntVarP(&opts.Cfg.Port, "port", "p", 5000, "Port the server will listen on.")
118-
115+
internal.ConfigFileFlags(flags, opts)
116+
internal.ServeFlags(flags, opts)
119117
flags.StringVar(&opts.ToolsFile, "tools_file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
120118
// deprecate tools_file
121119
_ = flags.MarkDeprecated("tools_file", "please use --tools-file instead")
122-
flags.BoolVar(&opts.Cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")
123120
flags.BoolVar(&opts.Cfg.DisableReload, "disable-reload", false, "Disables dynamic reloading of tools file.")
124-
flags.BoolVar(&opts.Cfg.UI, "ui", false, "Launches the Toolbox UI web server.")
125-
// TODO: Insecure by default. Might consider updating this for v1.0.0
126-
flags.StringSliceVar(&opts.Cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
127-
flags.StringSliceVar(&opts.Cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
128121
flags.IntVar(&opts.Cfg.PollInterval, "poll-interval", 0, "Specifies the polling frequency (seconds) for configuration file updates.")
129-
130122
// wrap RunE command so that we have access to original Command object
131123
cmd.RunE = func(*cobra.Command, []string) error { return run(cmd, opts) }
132124

133-
// Register subcommands for tool invocation
125+
// Register subcommands
134126
cmd.AddCommand(invoke.NewCommand(opts))
135-
// Register subcommands for skill generation
136127
cmd.AddCommand(skills.NewCommand(opts))
128+
cmd.AddCommand(serve.NewCommand(opts))
137129

138130
return cmd
139131
}

‎cmd/root_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -937,6 +937,7 @@ func TestSubcommandWiring(t *testing.T) {
937937
}{
938938
{[]string{"invoke"}, "invoke"},
939939
{[]string{"skills-generate"}, "skills-generate"},
940+
{[]string{"serve"}, "serve"},
940941
}
941942

942943
for _, tc := range tests {

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ require (
4848
github.com/sijms/go-ora/v2 v2.9.0
4949
github.com/snowflakedb/gosnowflake v1.19.0
5050
github.com/spf13/cobra v1.10.2
51+
github.com/spf13/pflag v1.0.9
5152
github.com/testcontainers/testcontainers-go v0.40.0
5253
github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0
5354
github.com/thlib/go-timezone-local v0.0.7
@@ -229,7 +230,6 @@ require (
229230
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
230231
github.com/shirou/gopsutil/v4 v4.25.6 // indirect
231232
github.com/sirupsen/logrus v1.9.3 // indirect
232-
github.com/spf13/pflag v1.0.9 // indirect
233233
github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect
234234
github.com/stretchr/testify v1.11.1 // indirect
235235
github.com/tklauser/go-sysconf v0.3.12 // indirect

0 commit comments

Comments
 (0)