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
10 changed files with 124 additions and 121 deletions

3
.gitignore vendored
View File

@@ -18,9 +18,6 @@ node_modules
# coverage
.coverage
# python
__pycache__/
# executable
genai-toolbox
toolbox

View File

@@ -1 +0,0 @@
GEMINI.md

View File

@@ -1 +0,0 @@
GEMINI.md

108
GEMINI.md
View File

@@ -1,108 +0,0 @@
# MCP Toolbox Context
This file (symlinked as `CLAUDE.md` and `AGENTS.md`) provides context and guidelines for AI agents working on the MCP Toolbox for Databases project. It summarizes key information from `CONTRIBUTING.md` and `DEVELOPER.md`.
## Project Overview
**MCP Toolbox for Databases** is a Go-based project designed to provide Model Context Protocol (MCP) tools for various data sources and services. It allows Large Language Models (LLMs) to interact with databases and other tools safely and efficiently.
## Tech Stack
- **Language:** Go (1.23+)
- **Documentation:** Hugo (Extended Edition v0.146.0+)
- **Containerization:** Docker
- **CI/CD:** GitHub Actions, Google Cloud Build
- **Linting:** `golangci-lint`
## Key Directories
- `cmd/`: Application entry points.
- `internal/sources/`: Implementations of database sources (e.g., Postgres, BigQuery).
- `internal/tools/`: Implementations of specific tools for each source.
- `tests/`: Integration tests.
- `docs/`: Project documentation (Hugo site).
## Development Workflow
### Prerequisites
- Go 1.23 or later.
- Docker (for building container images and running some tests).
- Access to necessary Google Cloud resources for integration testing (if applicable).
### Building and Running
1. **Build Binary:** `go build -o toolbox`
2. **Run Server:** `go run .` (Listens on port 5000 by default)
3. **Run with Help:** `go run . --help`
4. **Test Endpoint:** `curl http://127.0.0.1:5000`
### Testing
- **Unit Tests:** `go test -race -v ./cmd/... ./internal/...`
- **Integration Tests:**
- Run specific source tests: `go test -race -v ./tests/<source_dir>`
- Example: `go test -race -v ./tests/alloydbpg`
- Add new sources to `.ci/integration.cloudbuild.yaml`
- **Linting:** `golangci-lint run --fix`
## Developing Documentation
### Prerequisites
- Hugo (Extended Edition v0.146.0+)
- Node.js (for `npm ci`)
### Running Local Server
1. Navigate to `.hugo` directory: `cd .hugo`
2. Install dependencies: `npm ci`
3. Start server: `hugo server`
### Versioning Workflows
1. **Deploy In-development docs**: Merges to main -> `/dev/`.
2. **Deploy Versioned Docs**: New Release -> `/<version>/` and root.
3. **Deploy Previous Version Docs**: Manual workflow for older versions.
## Coding Conventions
### Tool Naming
- **Tool Name:** `snake_case` (e.g., `list_collections`, `run_query`).
- Do *not* include the product name (e.g., avoid `firestore_list_collections`).
- **Tool Type:** `kebab-case` (e.g., `firestore-list-collections`).
- *Must* include the product name.
### Branching and Commits
- **Branch Naming:** `feat/`, `fix/`, `docs/`, `chore/` (e.g., `feat/add-gemini-md`).
- **Commit Messages:** [Conventional Commits](https://www.conventionalcommits.org/) format.
- Format: `<type>(<scope>): <description>`
- Example: `feat(source/postgres): add new connection option`
- Types: `feat`, `fix`, `docs`, `chore`, `test`, `ci`, `refactor`, `revert`, `style`.
## Adding New Features
### Adding a New Data Source
1. Create a new directory: `internal/sources/<newdb>`.
2. Define `Config` and `Source` structs in `internal/sources/<newdb>/<newdb>.go`.
3. Implement `SourceConfig` interface (`SourceConfigType`, `Initialize`).
4. Implement `Source` interface (`SourceType`).
5. Implement `init()` to register the source.
6. Add unit tests in `internal/sources/<newdb>/<newdb>_test.go`.
### Adding a New Tool
1. Create a new directory: `internal/tools/<newdb>/<toolname>`.
2. Define `Config` and `Tool` structs.
3. Implement `ToolConfig` interface (`ToolConfigType`, `Initialize`).
4. Implement `Tool` interface (`Invoke`, `ParseParams`, `Manifest`, `McpManifest`, `Authorized`).
5. Implement `init()` to register the tool.
6. Add unit tests.
### Adding Documentation
- Add source documentation to `docs/en/resources/sources/`.
- Add tool documentation to `docs/en/resources/tools/`.

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(),
),
}