feat(server): allow partial source startup checks

Author: Deeven-Seru

cmd/internal/persistent_flags.go

@@ -43,4 +43,5 @@ func PersistentFlags(parentCmd *cobra.Command, opts *ToolboxOptions) {
 	)
 	persistentFlags.StringSliceVar(&opts.PrebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
 	persistentFlags.StringSliceVar(&opts.Cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
+	persistentFlags.BoolVar(&opts.Cfg.AllowPartialSources, "allow-partial-sources", false, "Skip startup connectivity checks for sources to allow partial availability.")
 }

cmd/internal/tools_file_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/googleapis/genai-toolbox/internal/prompts"
 	"github.com/googleapis/genai-toolbox/internal/prompts/custom"
 	"github.com/googleapis/genai-toolbox/internal/server"
+	"github.com/googleapis/genai-toolbox/internal/sources"
 	cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
 	httpsrc "github.com/googleapis/genai-toolbox/internal/sources/http"
 	"github.com/googleapis/genai-toolbox/internal/testutils"
@@ -2164,3 +2165,44 @@ tools:
 		})
 	}
 }
+
+func TestParseToolFileCheckAtStartup(t *testing.T) {
+	ctx, err := testutils.ContextWithNewLogger()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	input := testutils.FormatYaml(`
+	kind: sources
+	name: http-source
+	type: http
+	baseUrl: https://example.com
+	checkAtStartup: false
+	`)
+
+	toolsFile, err := parseToolsFile(ctx, input)
+	if err != nil {
+		t.Fatalf("failed to parse input: %v", err)
+	}
+	cfg, ok := toolsFile.Sources["http-source"]
+	if !ok {
+		t.Fatalf("expected source to be parsed")
+	}
+	if sources.CheckAtStartup(cfg) {
+		t.Fatalf("expected checkAtStartup to be false")
+	}
+
+	defaultInput := testutils.FormatYaml(`
+	kind: sources
+	name: http-source-default
+	type: http
+	baseUrl: https://example.com
+	`)
+	toolsFile, err = parseToolsFile(ctx, defaultInput)
+	if err != nil {
+		t.Fatalf("failed to parse input: %v", err)
+	}
+	cfg = toolsFile.Sources["http-source-default"]
+	if !sources.CheckAtStartup(cfg) {
+		t.Fatalf("expected checkAtStartup default to be true")
+	}
+}

cmd/root.go

@@ -138,13 +138,13 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
 	return cmd
 }
 
-func handleDynamicReload(ctx context.Context, toolsFile internal.ToolsFile, s *server.Server) error {
+func handleDynamicReload(ctx context.Context, toolsFile internal.ToolsFile, s *server.Server, allowPartialSources bool) error {
 	logger, err := util.LoggerFromContext(ctx)
 	if err != nil {
 		panic(err)
 	}
 
-	sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile)
+	sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile, allowPartialSources)
 	if err != nil {
 		errMsg := fmt.Errorf("unable to validate reloaded edits: %w", err)
 		logger.WarnContext(ctx, errMsg.Error())
@@ -158,7 +158,7 @@ func handleDynamicReload(ctx context.Context, toolsFile internal.ToolsFile, s *s
 
 // validateReloadEdits checks that the reloaded tools file configs can initialized without failing
 func validateReloadEdits(
-	ctx context.Context, toolsFile internal.ToolsFile,
+	ctx context.Context, toolsFile internal.ToolsFile, allowPartialSources bool,
 ) (map[string]sources.Source, map[string]auth.AuthService, map[string]embeddingmodels.EmbeddingModel, map[string]tools.Tool, map[string]tools.Toolset, map[string]prompts.Prompt, map[string]prompts.Promptset, error,
 ) {
 	logger, err := util.LoggerFromContext(ctx)
@@ -184,6 +184,7 @@ func validateReloadEdits(
 		ToolConfigs:           toolsFile.Tools,
 		ToolsetConfigs:        toolsFile.Toolsets,
 		PromptConfigs:         toolsFile.Prompts,
+		AllowPartialSources:   allowPartialSources,
 	}
 
 	sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, reloadedConfig)
@@ -239,7 +240,7 @@ func scanWatchedFiles(watchingFolder bool, folderToWatch string, watchedFiles ma
 }
 
 // watchChanges checks for changes in the provided yaml tools file(s) or folder.
-func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles map[string]bool, s *server.Server, pollTickerSecond int) {
+func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles map[string]bool, s *server.Server, pollTickerSecond int, allowPartialSources bool) {
 	logger, err := util.LoggerFromContext(ctx)
 	if err != nil {
 		panic(err)
@@ -384,7 +385,7 @@ func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles m
 				}
 			}
 
-			err = handleDynamicReload(ctx, reloadedToolsFile, s)
+			err = handleDynamicReload(ctx, reloadedToolsFile, s, allowPartialSources)
 			if err != nil {
 				errMsg := fmt.Errorf("unable to parse reloaded tools file at %q: %w", reloadedToolsFile, err)
 				logger.WarnContext(ctx, errMsg.Error())
@@ -500,7 +501,7 @@ func run(cmd *cobra.Command, opts *internal.ToolboxOptions) error {
 	if isCustomConfigured && !opts.Cfg.DisableReload {
 		watchDirs, watchedFiles := resolveWatcherInputs(opts.ToolsFile, opts.ToolsFiles, opts.ToolsFolder)
 		// start watching the file(s) or folder for changes to trigger dynamic reloading
-		go watchChanges(ctx, watchDirs, watchedFiles, s, opts.Cfg.PollInterval)
+		go watchChanges(ctx, watchDirs, watchedFiles, s, opts.Cfg.PollInterval, opts.Cfg.AllowPartialSources)
 	}
 
 	// wait for either the server to error out or the command's context to be canceled

docs/en/reference/cli.md

@@ -28,6 +28,7 @@ description: >
 |              | `--allowed-origins`        | Specifies a list of origins permitted to access this server for CORs access.                                                                                                     | `*`         |
 |              | `--allowed-hosts`          | Specifies a list of hosts permitted to access this server to prevent DNS rebinding attacks.                                                                                      | `*`         |
 |              | `--user-agent-metadata`    | Appends additional metadata to the User-Agent.                                                                                                                                   |             |
+|              | `--allow-partial-sources`  | Skip startup connectivity checks for sources to allow partial availability.                                                                                                      |             |
 |              | `--poll-interval`          | Specifies the polling frequency (seconds) for configuration file updates.                                                                                                        | `0`         |
 | `-v`         | `--version`                | version for toolbox                                                                                                                                                              |             |
 

docs/en/resources/sources/_index.md

@@ -26,9 +26,14 @@ instance: my-instance-name
 database: my_db
 user: ${USER_NAME}
 password: ${PASSWORD}
+# Optional: skip startup connectivity checks for this source
+checkAtStartup: false
 ```
 
 In implementation, each source is a different connection pool or client that used
 to connect to the database and execute the tool.
 
+By default, toolbox validates source connectivity at startup. Set `checkAtStartup: false`
+to allow a specific source to be offline without blocking startup.
+
 ## Available Sources

internal/server/config.go

@@ -41,6 +41,8 @@ type ServerConfig struct {
 	Port int
 	// SourceConfigs defines what sources of data are available for tools.
 	SourceConfigs SourceConfigs
+	// AllowPartialSources allows startup when some sources are unavailable by skipping connectivity checks.
+	AllowPartialSources bool
 	// AuthServiceConfigs defines what sources of authentication are available for tools.
 	AuthServiceConfigs AuthServiceConfigs
 	// EmbeddingModelConfigs defines a models used to embed parameters.
@@ -237,6 +239,17 @@ func UnmarshalYAMLSourceConfig(ctx context.Context, name string, r map[string]an
 	if !ok {
 		return nil, fmt.Errorf("missing 'type' field or it is not a string")
 	}
+
+	checkAtStartup := true
+	if raw, ok := r["checkAtStartup"]; ok {
+		val, ok := raw.(bool)
+		if !ok {
+			return nil, fmt.Errorf("checkAtStartup must be a boolean for source %q", name)
+		}
+		checkAtStartup = val
+		delete(r, "checkAtStartup")
+	}
+
 	dec, err := util.NewStrictDecoder(r)
 	if err != nil {
 		return nil, fmt.Errorf("error creating decoder: %w", err)
@@ -245,6 +258,9 @@ func UnmarshalYAMLSourceConfig(ctx context.Context, name string, r map[string]an
 	if err != nil {
 		return nil, err
 	}
+	if !checkAtStartup {
+		sourceConfig = sources.WrapSourceConfigWithStartupCheck(sourceConfig, checkAtStartup)
+	}
 	return sourceConfig, nil
 }
 

internal/server/server.go

@@ -81,7 +81,14 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
 
 	// initialize and validate the sources from configs
 	sourcesMap := make(map[string]sources.Source)
+	if cfg.AllowPartialSources {
+		l.WarnContext(ctx, "Skipping startup connectivity checks for all sources (allow-partial-sources enabled).")
+	}
 	for name, sc := range cfg.SourceConfigs {
+		checkAtStartup := sources.CheckAtStartup(sc)
+		if cfg.AllowPartialSources {
+			checkAtStartup = false
+		}
 		s, err := func() (sources.Source, error) {
 			childCtx, span := instrumentation.Tracer.Start(
 				ctx,
@@ -90,6 +97,10 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
 				trace.WithAttributes(attribute.String("source_name", name)),
 			)
 			defer span.End()
+			if !checkAtStartup && !cfg.AllowPartialSources {
+				l.InfoContext(childCtx, fmt.Sprintf("Skipping startup connectivity check for source %q", name))
+			}
+			childCtx = sources.WithStartupCheck(childCtx, checkAtStartup)
 			s, err := sc.Initialize(childCtx, instrumentation.Tracer)
 			if err != nil {
 				return nil, fmt.Errorf("unable to initialize source %q: %w", name, err)

internal/sources/alloydbpg/alloydb_pg.go

@@ -71,8 +71,9 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 		return nil, fmt.Errorf("unable to create pool: %w", err)
 	}
 
-	err = pool.Ping(ctx)
-	if err != nil {
+	if err := sources.CheckConnectivity(ctx, func(ctx context.Context) error {
+		return pool.Ping(ctx)
+	}); err != nil {
 		return nil, fmt.Errorf("unable to connect successfully: %w", err)
 	}
 

internal/sources/clickhouse/clickhouse.go

@@ -69,8 +69,9 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 		return nil, fmt.Errorf("unable to create pool: %w", err)
 	}
 
-	err = pool.PingContext(ctx)
-	if err != nil {
+	if err := sources.CheckConnectivity(ctx, func(ctx context.Context) error {
+		return pool.PingContext(ctx)
+	}); err != nil {
 		return nil, fmt.Errorf("unable to connect successfully: %w", err)
 	}
 

internal/sources/cloudsqlmssql/cloud_sql_mssql.go

@@ -75,8 +75,9 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 	}
 
 	// Verify db connection
-	err = db.PingContext(ctx)
-	if err != nil {
+	if err := sources.CheckConnectivity(ctx, func(ctx context.Context) error {
+		return db.PingContext(ctx)
+	}); err != nil {
 		return nil, fmt.Errorf("unable to connect successfully: %w", err)
 	}