Compare commits

...

4 Commits

Author SHA1 Message Date
Yuan Teoh
693616b2a9 resolve gemini coments 2026-02-13 11:12:48 -08:00
Yuan Teoh
af025b7834 feat: add polling system to dynamic reloading 2026-02-13 10:49:44 -08:00
Averi Kitsch
6602abd059 docs: update diagram (#2461)
## Description

> Should include a concise description of the changes (bug or feature),
it's
> impact, along with a summary of the solution

## PR Checklist

> Thank you for opening a Pull Request! Before submitting your PR, there
are a
> few things you can do to make sure it goes smoothly:

- [ ] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #<issue_number_goes_here>
2026-02-13 10:25:34 -08:00
Binh Tran
62b830987d fix: Deflake alloydb omni (#2431)
## Description

Improve logic to check that database is up.

**IMPORTANT** DO NOT MERGE until I have reverted
f7d7d9e708

## PR Checklist

- [x] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [x] https://github.com/googleapis/genai-toolbox/issues/2422
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #2422

Co-authored-by: Averi Kitsch <akitsch@google.com>
2026-02-13 09:55:53 -08:00
6 changed files with 124 additions and 8 deletions

View File

@@ -125,6 +125,7 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
// 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) }
@@ -195,8 +196,22 @@ func validateReloadEdits(
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
}
// Helper to check if a file has a newer ModTime than stored in the map
func checkModTime(path string, lastSeen map[string]time.Time) bool {
info, err := os.Stat(path)
if err != nil {
return false
}
mTime := info.ModTime()
if mTime.After(lastSeen[path]) {
lastSeen[path] = mTime
return true
}
return false
}
// 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) {
func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles map[string]bool, s *server.Server, pollTickerSecond int) {
logger, err := util.LoggerFromContext(ctx)
if err != nil {
panic(err)
@@ -238,6 +253,42 @@ func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles m
logger.DebugContext(ctx, fmt.Sprintf("Added directory %s to watcher.", dir))
}
lastSeen := make(map[string]time.Time)
var pollTickerChan <-chan time.Time
if pollTickerSecond > 0 {
ticker := time.NewTicker(time.Duration(pollTickerSecond) * time.Second)
defer ticker.Stop()
pollTickerChan = ticker.C // Assign the channel
logger.DebugContext(ctx, fmt.Sprintf("NFS polling enabled every %v", pollTickerSecond))
// Pre-populate lastSeen to avoid an initial spurious reload
if watchingFolder {
files, err := os.ReadDir(folderToWatch)
if err != nil {
logger.WarnContext(ctx, "error reading tools folder on initial scan %s", err)
} else {
for _, f := range files {
if !f.IsDir() && (strings.HasSuffix(f.Name(), ".yaml") || strings.HasSuffix(f.Name(), ".yml")) {
fullPath := filepath.Join(folderToWatch, f.Name())
info, err := os.Stat(fullPath)
if err == nil {
lastSeen[fullPath] = info.ModTime()
}
}
}
}
} else {
for f := range watchedFiles {
info, err := os.Stat(f)
if err == nil {
lastSeen[f] = info.ModTime()
}
}
}
} else {
logger.DebugContext(ctx, "NFS polling disabled (interval is 0)")
}
// debounce timer is used to prevent multiple writes triggering multiple reloads
debounceDelay := 100 * time.Millisecond
debounce := time.NewTimer(1 * time.Minute)
@@ -248,6 +299,55 @@ func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles m
case <-ctx.Done():
logger.DebugContext(ctx, "file watcher context cancelled")
return
case <-pollTickerChan:
changed := false
currentDiskFiles := make(map[string]bool)
// Get files that are currently on disk
if watchingFolder {
files, err := os.ReadDir(folderToWatch)
if err != nil {
logger.WarnContext(ctx, "error reading tools folder %s", err)
continue
}
for _, f := range files {
if !f.IsDir() && (strings.HasSuffix(f.Name(), ".yaml") || strings.HasSuffix(f.Name(), ".yml")) {
fullPath := filepath.Join(folderToWatch, f.Name())
currentDiskFiles[fullPath] = true
if checkModTime(fullPath, lastSeen) {
changed = true
}
}
}
} else {
for f := range watchedFiles {
// We must explicitly check existence here because checkModTime
// swallows errors (returns false), making it impossible to
// distinguish between "no change" and "file deleted".
if _, err := os.Stat(f); err == nil {
currentDiskFiles[f] = true
if checkModTime(f, lastSeen) {
changed = true
}
}
}
}
// Check for Deletions
// If it was in lastSeen but is NOT in currentDiskFiles, it's
// deleted; we will need to reload the server.
for path := range lastSeen {
if !currentDiskFiles[path] {
logger.DebugContext(ctx, fmt.Sprintf("File deleted (detected via polling): %s", path))
delete(lastSeen, path)
changed = true
}
}
if changed {
logger.DebugContext(ctx, "NFS remote change detected via polling")
// once this timer runs out, it will trigger debounce.C
debounce.Reset(debounceDelay)
}
case err, ok := <-w.Errors:
if !ok {
logger.WarnContext(ctx, "file watcher was closed unexpectedly")
@@ -417,7 +517,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)
go watchChanges(ctx, watchDirs, watchedFiles, s, opts.Cfg.PollInterval)
}
// wait for either the server to error out or the command's context to be canceled

View File

@@ -590,7 +590,7 @@ func TestSingleEdit(t *testing.T) {
watchedFiles := map[string]bool{cleanFileToWatch: true}
watchDirs := map[string]bool{watchDir: true}
go watchChanges(ctx, watchDirs, watchedFiles, mockServer)
go watchChanges(ctx, watchDirs, watchedFiles, mockServer, 0)
// escape backslash so regex doesn't fail on windows filepaths
regexEscapedPathFile := strings.ReplaceAll(cleanFileToWatch, `\`, `\\\\*\\`)

Binary file not shown.

Before

Width:  |  Height:  |  Size: 241 KiB

After

Width:  |  Height:  |  Size: 271 KiB

View File

@@ -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. | |
| | `--poll-interval` | Specifies the polling frequency (seconds) for configuration file updates. | `0` |
| `-v` | `--version` | version for toolbox | |
## Sub Commands
@@ -136,6 +137,16 @@ used at a time.
Toolbox enables dynamic reloading by default. To disable, use the
`--disable-reload` flag.
Use the `--poll-interval` flag to manually detect configuration file updates.
When the poll interval is `0`, the polling system is disabled.
{{< notice tip >}}
For polling system to be effective when running Kubernetes, PersistentVolume or
StorageClass must be set to refresh attributes rapidly by setting `actimeo=1`.
Actimeo setting determines the duration for which a client trusts its local
cache for file attributes.
{{< /notice >}}
### Toolbox UI
To launch Toolbox's interactive UI, use the `--ui` flag. This allows you to test

View File

@@ -75,6 +75,8 @@ type ServerConfig struct {
AllowedHosts []string
// UserAgentMetadata specifies additional metadata to append to the User-Agent string.
UserAgentMetadata []string
// PollInterval sets the polling frequency for configuration file updates.
PollInterval int
}
type logFormat string

View File

@@ -36,15 +36,17 @@ var (
AlloyDBDatabase = "postgres"
)
// Copied over from postgres.go
func initPostgresConnectionPool(host, port, user, pass, dbname string) (*pgxpool.Pool, error) {
// urlExample := "postgres:dd//username:password@localhost:5432/database_name"
url := &url.URL{
func buildPostgresURL(host, port, user, pass, dbname string) *url.URL {
return &url.URL{
Scheme: "postgres",
User: url.UserPassword(user, pass),
Host: fmt.Sprintf("%s:%s", host, port),
Path: dbname,
}
}
func initPostgresConnectionPool(host, port, user, pass, dbname string) (*pgxpool.Pool, error) {
url := buildPostgresURL(host, port, user, pass, dbname)
pool, err := pgxpool.New(context.Background(), url.String())
if err != nil {
return nil, fmt.Errorf("Unable to create connection pool: %w", err)
@@ -63,7 +65,8 @@ func setupAlloyDBContainer(ctx context.Context, t *testing.T) (string, string, f
"POSTGRES_PASSWORD": AlloyDBPass,
},
WaitingFor: wait.ForAll(
wait.ForLog("Post Startup: Successfully reinstalled extensions"),
wait.ForLog("database system was shut down at"),
wait.ForLog("database system is ready to accept connections"),
wait.ForExposedPort(),
),
}