Skip to content

Commit 7fdbef9

Browse files
akangsha7Akangsha Goel
andauthored
fix: revert "feat(sources): connect sources on first use" (#4137)
## Description Reverts #4076 (`5700630`) to unbreak v1.13.0 while the fix is worked through separately. #4076 moved source setup into `ConnectOnce.Do`, which runs the connect under a context it cancels as soon as the connect returns. An oauth2 token source keeps the context it was built with and reuses it for every later refresh, so a client built there fails its first refresh and never recovers: ``` error processing GCP request: Get ".../tables/?alt=json&prettyPrint=false": Post "https://oauth2.googleapis.com/token": context canceled ``` Reported against Data Agent Kit. Every source was migrated in #4076, so the same shape reaches `alloydbadmin`, `cloudsqladmin`, `cloudmonitoring`, `databaseinsights`, `cloudhealthcare`, `cloudloggingadmin`, `datalineage`, `dataplex`, `looker`, plus `bigtable` and `valkey` through a different mechanism. Details in #4135. Reverting restores connecting at startup, under a context that lives as long as the server, which is where the credentials were built before v1.13.0. Titled `fix:` rather than `revert:` deliberately, so release-please cuts a patch release. ### What is kept The `sources.ConnectOnce` helper from #3905 stays in place, unused. Only the per-source wiring and the `--defer-source-connect` flag are reverted, so the feature can be re-landed without rebuilding the helper. ### Re-land #4134 has the actual fix — a `DetachedConnectContext` for anything built inside a connect that outlives it — along with a regression test that fails on `5700630` and on v1.13.0, and passes with the fix. It should be rebased on top of this revert and re-landed together with #4076. ## Verification Clean revert, no conflicts. On this branch: ``` go build ./... go test -race ./cmd/... ./internal/... golangci-lint run ``` all pass. ## PR Checklist - [x] Make sure to open an issue as a bug/issue before writing your code! - [x] Ensure you have manually reviewed the entire diff before requesting a review - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involves a breaking change ## Issue Reference Fixes#4135 🦕 Co-authored-by: Akangsha Goel
1 parent 8896373 commit 7fdbef9

195 files changed

Lines changed: 1950 additions & 3347 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎cmd/internal/flags.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,5 +79,4 @@ func ServeFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
7979
flags.BoolVar(&opts.Cfg.EnableDraftSpecs, "enable-draft-specs", false, "Opt-in and test upcoming draft MCP specifications.")
8080
flags.StringSliceVar(&opts.Cfg.DisableExt, "disable-ext", []string{}, "Specifies MCP extension URIs disabled on this server.")
8181
flags.StringVar(&opts.Cfg.OpenAIAppsChallengeFile, "openai-apps-challenge-file", "", "Path to a file containing the OpenAI verification challenge token to serve at /.well-known/openai-apps-challenge.")
82-
flags.BoolVar(&opts.Cfg.DeferSourceConnect, "defer-source-connect", false, "Connect to each source on first use instead of at startup. Tools can be listed without any source being reachable.")
8382
}

‎cmd/root_test.go‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -304,13 +304,6 @@ func TestServerConfigFlags(t *testing.T) {
304304
OpenAIAppsChallengeFile: "openai-token.txt",
305305
}),
306306
},
307-
{
308-
desc: "defer source connect",
309-
args: []string{"--defer-source-connect"},
310-
want: withDefaults(server.ServerConfig{
311-
DeferSourceConnect: true,
312-
}),
313-
},
314307
}
315308
for _, tc := range tcs {
316309
t.Run(tc.desc, func(t *testing.T) {

‎docs/en/reference/cli.md‎

Lines changed: 0 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ description: >
1111
| Flag (Short) | Flag (Long) | Description | Default |
1212
|--------------|----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------|
1313
| `-a` | `--address` | Address of the interface the server will listen on. | `127.0.0.1` |
14-
| | `--defer-source-connect` | Connect to each source on first use instead of at startup. | `false` |
1514
| | `--disable-ext` | Specifies MCP extension URIs disabled on this server. | |
1615
| | `--disable-reload` | Disables dynamic reloading config. | |
1716
| `-h` | `--help` | help for toolbox | |
@@ -246,37 +245,6 @@ To launch Toolbox's interactive UI, use the `--ui` flag. This allows you to test
246245
tools and toolsets with features such as authorized parameters. To learn more,
247246
visit [Toolbox UI](../documentation/configuration/toolbox-ui/index.md).
248247

249-
### Deferring Source Connections
250-
251-
By default Toolbox connects to every configured source at startup and fails to
252-
start if any connection fails. Pass `--defer-source-connect` to defer each
253-
connection until the first tool call that needs it.
254-
255-
```bash
256-
./toolbox --tools-file tools.yaml --defer-source-connect
257-
```
258-
259-
This is useful when you want to inspect a tool catalog without provisioning
260-
databases, when cold start time matters, or when you would rather a broken
261-
source surface as a tool error the agent can read than as a server that never
262-
comes up. With the flag set:
263-
264-
* Toolbox starts even when no source is reachable, and `tools/list` and
265-
`/api/toolset` return the full catalog with complete schemas.
266-
* The first call to a tool connects its source. If that fails, the call returns
267-
a tool error containing the connection failure and the server stays up. Only
268-
successful connections are kept, so a source that comes up later starts
269-
working without a restart.
270-
* A tool naming a source that does not exist, or one whose type it cannot use,
271-
is still rejected at startup.
272-
273-
Configuration is still validated at startup. A malformed value — an unparseable
274-
`queryTimeout`, an invalid `writeMode` — fails immediately, because detecting it
275-
takes no network.
276-
277-
Environment variables are still required, since a source cannot connect without
278-
them, and an unset one fails startup as it does today.
279-
280248
### Disabling MCP Extensions
281249

282250
By default, Toolbox advertises support for its own custom MCP extensions (e.g., `com.google.cloud/toolbox.v1`) during the client discovery phase. This extension signals to clients that they can leverage Toolbox-specific features that fall outside the official MCP specification (see the [Extension README](https://github.com/googleapis/mcp-toolbox/blob/main/extensions/2026-07-28/README.md) for a list of currently supported capabilities).

‎internal/server/config.go‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,6 @@ type ServerConfig struct {
114114
DisableVersionCheck bool
115115
// OpenAIAppsChallengeFile specifies the path to a file containing the OpenAI verification challenge token to serve at /.well-known/openai-apps-challenge.
116116
OpenAIAppsChallengeFile string
117-
// DeferSourceConnect connects each source on first use instead of at startup.
118-
DeferSourceConnect bool
119117
}
120118

121119
type logFormat string

‎internal/server/primitives/primitives_test.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,17 @@ import (
2525
"github.com/googleapis/mcp-toolbox/internal/resources"
2626
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
2727
"github.com/googleapis/mcp-toolbox/internal/sources"
28+
"github.com/googleapis/mcp-toolbox/internal/sources/alloydbpg"
2829
"github.com/googleapis/mcp-toolbox/internal/testutils"
2930
"github.com/googleapis/mcp-toolbox/internal/tools"
3031
)
3132

3233
func TestUpdateServer(t *testing.T) {
3334
newSources := map[string]sources.Source{
34-
"example-source": testutils.MockSource{
35-
MockSourceConfig: testutils.MockSourceConfig{
36-
Name: "example-source",
37-
Type: "mock-source",
35+
"example-source": &alloydbpg.Source{
36+
Config: alloydbpg.Config{
37+
Name: "example-alloydb-source",
38+
Type: "alloydb-postgres",
3839
},
3940
},
4041
}
@@ -88,10 +89,10 @@ func TestUpdateServer(t *testing.T) {
8889
t.Errorf("error updating server, resource templates (-want +got):\n%s", diff)
8990
}
9091
updateSource := map[string]sources.Source{
91-
"example-source2": testutils.MockSource{
92-
MockSourceConfig: testutils.MockSourceConfig{
93-
Name: "example-source2",
94-
Type: "mock-source",
92+
"example-source2": &alloydbpg.Source{
93+
Config: alloydbpg.Config{
94+
Name: "example-alloydb-source2",
95+
Type: "alloydb-postgres",
9596
},
9697
},
9798
}

‎internal/server/server.go‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
117117
trace.WithAttributes(attribute.String("source_name", name)),
118118
)
119119
defer span.End()
120-
s, err := sc.Initialize(childCtx, instrumentation.Tracer, cfg.DeferSourceConnect)
120+
s, err := sc.Initialize(childCtx, instrumentation.Tracer)
121121
if err != nil {
122122
return nil, fmt.Errorf("unable to initialize source %q: %w", name, err)
123123
}
@@ -133,9 +133,6 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
133133
sourceNames = append(sourceNames, name)
134134
}
135135
l.InfoContext(ctx, fmt.Sprintf("Initialized %d sources: %s", len(sourcesMap), strings.Join(sourceNames, ", ")))
136-
if cfg.DeferSourceConnect {
137-
l.InfoContext(ctx, "Source connections are deferred; each source connects on first use.")
138-
}
139136

140137
// initialize and validate the auth services from configs
141138
authServicesMap := make(map[string]auth.AuthService)

‎internal/server/server_test.go‎

Lines changed: 6 additions & 166 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import (
2424
"crypto/x509/pkix"
2525
"encoding/json"
2626
"encoding/pem"
27-
"errors"
2827
"fmt"
2928
"io"
3029
"math/big"
@@ -35,7 +34,6 @@ import (
3534
"reflect"
3635
"slices"
3736
"strings"
38-
"sync/atomic"
3937
"testing"
4038
"time"
4139

@@ -54,7 +52,7 @@ import (
5452
"github.com/googleapis/mcp-toolbox/internal/server"
5553
v20260728 "github.com/googleapis/mcp-toolbox/internal/server/mcp/v20260728"
5654
"github.com/googleapis/mcp-toolbox/internal/sources"
57-
_ "github.com/googleapis/mcp-toolbox/internal/sources/alloydbpg"
55+
"github.com/googleapis/mcp-toolbox/internal/sources/alloydbpg"
5856
_ "github.com/googleapis/mcp-toolbox/internal/sources/postgres"
5957
_ "github.com/googleapis/mcp-toolbox/internal/sources/sqlite"
6058
"github.com/googleapis/mcp-toolbox/internal/telemetry"
@@ -64,7 +62,6 @@ import (
6462
"github.com/googleapis/mcp-toolbox/internal/tools/mysql/mysqlexecutesql"
6563
"github.com/googleapis/mcp-toolbox/internal/tools/postgres/postgresexecutesql"
6664
"github.com/googleapis/mcp-toolbox/internal/util"
67-
"go.opentelemetry.io/otel/trace"
6865
)
6966

7067
// Helper function to create temporary self-signed certs for the test
@@ -418,10 +415,10 @@ func TestUpdateServer(t *testing.T) {
418415
}
419416

420417
newSources := map[string]sources.Source{
421-
"example-source": testutils.MockSource{
422-
MockSourceConfig: testutils.MockSourceConfig{
423-
Name: "example-source",
424-
Type: "mock-source",
418+
"example-source": &alloydbpg.Source{
419+
Config: alloydbpg.Config{
420+
Name: "example-alloydb-source",
421+
Type: "alloydb-postgres",
425422
},
426423
},
427424
}
@@ -1515,7 +1512,7 @@ func TestInitializeConfigs(t *testing.T) {
15151512
ctx = util.WithInstrumentation(ctx, instrumentation)
15161513
t.Run("valid initialization", func(t *testing.T) {
15171514
sourceConfig1 := testutils.MockSourceConfig{Name: "my-source", Type: "mock-source"}
1518-
source1, _ := sourceConfig1.Initialize(ctx, nil, false)
1515+
source1, _ := sourceConfig1.Initialize(ctx, nil)
15191516
tools1 := testutils.NewMockTool("my-tool", "mock tool for offline config", "my-source", nil, false, false)
15201517
validCfg := server.ServerConfig{
15211518
Version: "0.0.0",
@@ -2275,160 +2272,3 @@ func TestOpenAIAppsChallenge(t *testing.T) {
22752272
}
22762273
})
22772274
}
2278-
2279-
var errSourceUnreachable = errors.New("source is unreachable")
2280-
2281-
// countingSourceConfig mirrors the shape every real source implements: it
2282-
// returns its Source without connecting when deferConnect is set, and resolves
2283-
// the handle through sources.ConnectOnce on first use.
2284-
type countingSourceConfig struct {
2285-
name string
2286-
connectErr error
2287-
connects *atomic.Int32
2288-
}
2289-
2290-
func (c countingSourceConfig) SourceConfigType() string { return "counting-source" }
2291-
2292-
func (c countingSourceConfig) Initialize(ctx context.Context, tracer trace.Tracer, deferConnect bool) (sources.Source, error) {
2293-
s := &countingSource{
2294-
cfg: c,
2295-
conn: sources.NewConnectOnce[*int](ctx, c.name, "counting-source", tracer),
2296-
}
2297-
if deferConnect {
2298-
return s, nil
2299-
}
2300-
if _, err := s.handle(ctx); err != nil {
2301-
return nil, err
2302-
}
2303-
return s, nil
2304-
}
2305-
2306-
type countingSource struct {
2307-
cfg countingSourceConfig
2308-
conn *sources.ConnectOnce[*int]
2309-
}
2310-
2311-
func (s *countingSource) handle(ctx context.Context) (*int, error) {
2312-
return s.conn.Do(ctx, func(ctx context.Context) (*int, error) {
2313-
s.cfg.connects.Add(1)
2314-
if s.cfg.connectErr != nil {
2315-
return nil, s.cfg.connectErr
2316-
}
2317-
return new(int), nil
2318-
})
2319-
}
2320-
2321-
func (s *countingSource) SourceType() string { return "counting-source" }
2322-
func (s *countingSource) ToConfig() sources.SourceConfig { return s.cfg }
2323-
func (s *countingSource) IsReadOnly() bool { return false }
2324-
2325-
func TestInitializeConfigsDeferSourceConnect(t *testing.T) {
2326-
ctx, err := testutils.ContextWithNewLogger()
2327-
if err != nil {
2328-
t.Fatalf("error setting up logger: %s", err)
2329-
}
2330-
instrumentation, err := telemetry.CreateTelemetryInstrumentation("0.0.0")
2331-
if err != nil {
2332-
t.Fatalf("unexpected error: %s", err)
2333-
}
2334-
ctx = util.WithInstrumentation(ctx, instrumentation)
2335-
2336-
newCfg := func(deferConnect bool, connects *atomic.Int32, connectErr error) server.ServerConfig {
2337-
return server.ServerConfig{
2338-
Version: "0.0.0",
2339-
SourceConfigs: server.SourceConfigs{
2340-
"my-source": countingSourceConfig{name: "my-source", connectErr: connectErr, connects: connects},
2341-
},
2342-
DeferSourceConnect: deferConnect,
2343-
SkipSourceValidation: true,
2344-
}
2345-
}
2346-
2347-
tcs := []struct {
2348-
desc string
2349-
deferConnect bool
2350-
connectErr error
2351-
wantErr bool
2352-
wantConnects int32
2353-
}{
2354-
{
2355-
desc: "flag off connects at startup",
2356-
deferConnect: false,
2357-
wantConnects: 1,
2358-
},
2359-
{
2360-
desc: "flag off fails startup when the source is unreachable",
2361-
deferConnect: false,
2362-
connectErr: errSourceUnreachable,
2363-
wantErr: true,
2364-
wantConnects: 1,
2365-
},
2366-
{
2367-
desc: "flag on skips connecting at startup",
2368-
deferConnect: true,
2369-
wantConnects: 0,
2370-
},
2371-
{
2372-
desc: "flag on starts up even when the source is unreachable",
2373-
deferConnect: true,
2374-
connectErr: errSourceUnreachable,
2375-
wantConnects: 0,
2376-
},
2377-
}
2378-
for _, tc := range tcs {
2379-
t.Run(tc.desc, func(t *testing.T) {
2380-
var connects atomic.Int32
2381-
sourcesMap, _, _, _, _, _, _, _, err := server.InitializeConfigs(ctx, newCfg(tc.deferConnect, &connects, tc.connectErr))
2382-
switch {
2383-
case tc.wantErr && err == nil:
2384-
t.Fatalf("expected an error but got nil")
2385-
case tc.wantErr && !errors.Is(err, tc.connectErr):
2386-
t.Fatalf("expected the connect failure to surface, got %q", err)
2387-
case !tc.wantErr && err != nil:
2388-
t.Fatalf("unexpected error: %s", err)
2389-
}
2390-
if got := connects.Load(); got != tc.wantConnects {
2391-
t.Fatalf("connect attempts during startup: got %d, want %d", got, tc.wantConnects)
2392-
}
2393-
if tc.wantErr {
2394-
return
2395-
}
2396-
// A deferred source is still listed, so its tools stay invocable.
2397-
if _, ok := sourcesMap["my-source"]; !ok {
2398-
t.Fatalf("expected %q in the sources map, got %v", "my-source", sourcesMap)
2399-
}
2400-
})
2401-
}
2402-
2403-
t.Run("flag on connects once on first use", func(t *testing.T) {
2404-
var connects atomic.Int32
2405-
sourcesMap, _, _, _, _, _, _, _, err := server.InitializeConfigs(ctx, newCfg(true, &connects, nil))
2406-
if err != nil {
2407-
t.Fatalf("unexpected error: %s", err)
2408-
}
2409-
s := sourcesMap["my-source"].(*countingSource)
2410-
for i := range 3 {
2411-
if _, err := s.handle(ctx); err != nil {
2412-
t.Fatalf("use %d failed: %s", i, err)
2413-
}
2414-
}
2415-
if got := connects.Load(); got != 1 {
2416-
t.Fatalf("connect attempts across three uses: got %d, want 1", got)
2417-
}
2418-
})
2419-
2420-
t.Run("flag on surfaces the connect failure at first use", func(t *testing.T) {
2421-
var connects atomic.Int32
2422-
sourcesMap, _, _, _, _, _, _, _, err := server.InitializeConfigs(ctx, newCfg(true, &connects, errSourceUnreachable))
2423-
if err != nil {
2424-
t.Fatalf("unexpected error: %s", err)
2425-
}
2426-
s := sourcesMap["my-source"].(*countingSource)
2427-
if _, err := s.handle(ctx); !errors.Is(err, errSourceUnreachable) {
2428-
t.Fatalf("expected the connect failure at first use, got %q", err)
2429-
}
2430-
if got := connects.Load(); got != 1 {
2431-
t.Fatalf("connect attempts: got %d, want 1", got)
2432-
}
2433-
})
2434-
}

0 commit comments

Comments
 (0)