Skip to content

Commit 3f04423

Browse files
Yuan325Yuan325
committed
feat: add serve subcommand
1 parent 61d0ea4 commit 3f04423

File tree

8 files changed

+231
-22
lines changed

8 files changed

+231
-22
lines changed

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

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,27 +20,40 @@ 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+
func ConfigFileFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
40+
flags.StringVar(&opts.ToolsFile, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
41+
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.")
42+
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.")
3943
// Fetch prebuilt tools sources to customize the help description
4044
prebuiltHelp := fmt.Sprintf(
4145
"Use a prebuilt tool configuration by source type. Allowed: '%s'. Can be specified multiple times.",
4246
strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"),
4347
)
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.")
48+
flags.StringSliceVar(&opts.PrebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
49+
}
50+
51+
func ServeFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
52+
flags.StringVarP(&opts.Cfg.Address, "address", "a", "127.0.0.1", "Address of the interface the server will listen on.")
53+
flags.IntVarP(&opts.Cfg.Port, "port", "p", 5000, "Port the server will listen on.")
54+
flags.BoolVar(&opts.Cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")
55+
flags.BoolVar(&opts.Cfg.UI, "ui", false, "Launches the Toolbox UI web server.")
56+
57+
flags.StringSliceVar(&opts.Cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
58+
flags.StringSliceVar(&opts.Cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
4659
}

‎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, "Sending 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)
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: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
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+
"fmt"
21+
"strings"
22+
"testing"
23+
"time"
24+
25+
"github.com/googleapis/genai-toolbox/cmd/internal"
26+
"github.com/spf13/cobra"
27+
)
28+
29+
func serveCommand(ctx context.Context, args []string) (string, error) {
30+
parentCmd := &cobra.Command{Use: "toolbox"}
31+
32+
buf := new(bytes.Buffer)
33+
opts := internal.NewToolboxOptions(internal.WithIOStreams(buf, buf))
34+
internal.PersistentFlags(parentCmd, opts)
35+
36+
cmd := NewCommand(opts)
37+
parentCmd.AddCommand(cmd)
38+
parentCmd.SetArgs(args)
39+
// Inject the context into the Cobra command
40+
parentCmd.SetContext(ctx)
41+
42+
err := parentCmd.Execute()
43+
return buf.String(), err
44+
}
45+
46+
func TestServe(t *testing.T) {
47+
fmt.Println("here")
48+
// context will automatically shutdown in 1 second.
49+
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
50+
defer cancel()
51+
52+
args := []string{"serve", "--port", "0"}
53+
output, err := serveCommand(ctx, args)
54+
if err != nil {
55+
t.Fatalf("expected graceful shutdown without error, got: %v", err)
56+
}
57+
58+
if !strings.Contains(output, "Server ready to serve!") {
59+
t.Errorf("expected to find server ready message in output, got: %s", output)
60+
}
61+
62+
if !strings.Contains(output, "Shutting down gracefully...") {
63+
t.Errorf("expected to find graceful shutdown message in output, got: %s", output)
64+
}
65+
}

‎cmd/internal/skills/command.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,12 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
5050
},
5151
}
5252

53-
cmd.Flags().StringVar(&cmd.name, "name", "", "Name of the generated skill.")
54-
cmd.Flags().StringVar(&cmd.description, "description", "", "Description of the generated skill")
55-
cmd.Flags().StringVar(&cmd.toolset, "toolset", "", "Name of the toolset to convert into a skill. If not provided, all tools will be included.")
56-
cmd.Flags().StringVar(&cmd.outputDir, "output-dir", "skills", "Directory to output generated skills")
53+
flags := cmd.Flags()
54+
internal.ConfigFileFlags(flags, opts)
55+
flags.StringVar(&cmd.name, "name", "", "Name of the generated skill.")
56+
flags.StringVar(&cmd.description, "description", "", "Description of the generated skill")
57+
flags.StringVar(&cmd.toolset, "toolset", "", "Name of the toolset to convert into a skill. If not provided, all tools will be included.")
58+
flags.StringVar(&cmd.outputDir, "output-dir", "skills", "Directory to output generated skills")
5759

5860
_ = cmd.MarkFlagRequired("name")
5961
_ = cmd.MarkFlagRequired("description")

‎cmd/root.go‎

Lines changed: 5 additions & 11 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,29 +111,22 @@ 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.")
125121
// 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 '*'.")
128122

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

132-
// Register subcommands for tool invocation
126+
// Register subcommands
133127
cmd.AddCommand(invoke.NewCommand(opts))
134-
// Register subcommands for skill generation
135128
cmd.AddCommand(skills.NewCommand(opts))
129+
cmd.AddCommand(serve.NewCommand(opts))
136130

137131
return cmd
138132
}

‎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"}, "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.18.1
5050
github.com/spf13/cobra v1.10.1
51+
github.com/spf13/pflag v1.0.9
5152
github.com/testcontainers/testcontainers-go v0.40.0
5253
github.com/thlib/go-timezone-local v0.0.7
5354
github.com/trinodb/trino-go-client v0.330.0
@@ -228,7 +229,6 @@ require (
228229
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
229230
github.com/shirou/gopsutil/v4 v4.25.6 // indirect
230231
github.com/sirupsen/logrus v1.9.3 // indirect
231-
github.com/spf13/pflag v1.0.9 // indirect
232232
github.com/spiffe/go-spiffe/v2 v2.5.0 // indirect
233233
github.com/stretchr/testify v1.11.1 // indirect
234234
github.com/tklauser/go-sysconf v0.3.12 // indirect

0 commit comments

Comments
 (0)