Skip to content

Commit 016245c

Browse files
duwenxin99duwenxin99
authored
feat(server/mcp): introduce generic client extension registry (#3723)
1. Description - Added mcp.InitializeProtocols to orchestrate startup setup in internal/server/mcp, decoupling server.go from protocol-specific packages. - Add generic extension infra to the 20260728 mcp version. 2. PR Checklist - [x] Make sure to open an issue as a bug/issue before writing your code! - [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 3. Issue Reference Fixes #0 🦕
1 parent 0adeaa5 commit 016245c

10 files changed

Lines changed: 276 additions & 3 deletions

File tree

‎cmd/internal/flags.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,4 +76,5 @@ func ServeFlags(flags *pflag.FlagSet, opts *ToolboxOptions) {
7676
flags.StringSliceVar(&opts.Cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
7777
flags.Int64Var(&opts.Cfg.HttpMaxRequestBytes, "http-max-request-bytes", server.DefaultHTTPMaxRequestBytes, "Maximum MCP HTTP request body size in bytes.")
7878
flags.BoolVar(&opts.Cfg.EnableDraftSpecs, "enable-draft-specs", false, "Opt-in and test upcoming draft MCP specifications.")
79+
flags.StringSliceVar(&opts.Cfg.DisableExt, "disable-ext", []string{}, "Specifies MCP extension URIs disabled on this server.")
7980
}

‎cmd/root_test.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,9 @@ func withDefaults(c server.ServerConfig) server.ServerConfig {
6363
if c.UserAgentMetadata == nil {
6464
c.UserAgentMetadata = []string{}
6565
}
66+
if c.DisableExt == nil {
67+
c.DisableExt = []string{}
68+
}
6669
if c.HttpMaxRequestBytes == 0 {
6770
c.HttpMaxRequestBytes = server.DefaultHTTPMaxRequestBytes
6871
}
@@ -256,6 +259,13 @@ func TestServerConfigFlags(t *testing.T) {
256259
KeyFile: "key.pem",
257260
}),
258261
},
262+
{
263+
desc: "disable ext",
264+
args: []string{"--disable-ext", "io.modelcontextprotocol/tasks"},
265+
want: withDefaults(server.ServerConfig{
266+
DisableExt: []string{"io.modelcontextprotocol/tasks"},
267+
}),
268+
},
259269
}
260270
for _, tc := range tcs {
261271
t.Run(tc.desc, func(t *testing.T) {

‎internal/server/config.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,8 @@ type ServerConfig struct {
103103
EnableDraftSpecs bool
104104
// SkipSourceValidation skips source validation during server startup
105105
SkipSourceValidation bool
106+
// DisableExt specifies MCP extension URIs disabled on this server.
107+
DisableExt []string
106108
}
107109

108110
type logFormat string

‎internal/server/mcp/mcp.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,16 @@ import (
3232
"github.com/googleapis/mcp-toolbox/internal/util"
3333
)
3434

35+
// ProtocolOptions contains configuration passed during server initialization to protocol handlers.
36+
type ProtocolOptions struct {
37+
DisableExt []string
38+
}
39+
40+
// InitializeProtocols performs version-specific protocol setup across all supported MCP versions.
41+
func InitializeProtocols(opts ProtocolOptions) {
42+
v20260728.Initialize(opts.DisableExt)
43+
}
44+
3545
// NotificationHandler process notifications request. It MUST NOT send a response.
3646
// Currently Toolbox does not process any notifications.
3747
func NotificationHandler(ctx context.Context, body []byte) error {

‎internal/server/mcp/v20260728/extensions.go‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
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 v20260728
16+
17+
import "slices"
18+
19+
// SupportedExtensions lists all MCP extension URIs supported by Toolbox by default.
20+
var SupportedExtensions = map[string]any{}
21+
22+
// ServerExtensions is the map of extension URIs enabled on this server.
23+
var ServerExtensions map[string]any
24+
25+
// Initialize performs version-specific protocol setup for v20260728.
26+
func Initialize(disabledExts []string) {
27+
ServerExtensions = make(map[string]any)
28+
for ext, extConfig := range SupportedExtensions {
29+
if ext != "" && !slices.Contains(disabledExts, ext) {
30+
ServerExtensions[ext] = extConfig
31+
}
32+
}
33+
}
34+
35+
// ParseSupportedExtensions returns a map of extension URIs that are supported by both the client and the server.
36+
func ParseSupportedExtensions(clientExtensions map[string]any) map[string]any {
37+
supported := make(map[string]any)
38+
if len(clientExtensions) == 0 || len(ServerExtensions) == 0 {
39+
return supported
40+
}
41+
for uri, clientExtVal := range clientExtensions {
42+
if _, ok := ServerExtensions[uri]; ok {
43+
supported[uri] = clientExtVal
44+
}
45+
}
46+
return supported
47+
}

‎internal/server/mcp/v20260728/extensions_test.go‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
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 v20260728
16+
17+
import (
18+
"testing"
19+
)
20+
21+
const testExtURI = "com.google.cloud/test-extension"
22+
23+
func TestParseSupportedExtensions(t *testing.T) {
24+
orig := ServerExtensions
25+
t.Cleanup(func() {
26+
ServerExtensions = orig
27+
})
28+
tests := []struct {
29+
name string
30+
extensions map[string]any
31+
serverExts map[string]any
32+
expectedUri string
33+
expectedVal bool
34+
}{
35+
{
36+
name: "nil map",
37+
extensions: nil,
38+
serverExts: nil,
39+
expectedUri: testExtURI,
40+
expectedVal: false,
41+
},
42+
{
43+
name: "enabled extension via empty settings object in extensions",
44+
extensions: map[string]any{
45+
testExtURI: map[string]any{},
46+
},
47+
serverExts: nil,
48+
expectedUri: testExtURI,
49+
expectedVal: true,
50+
},
51+
{
52+
name: "enabled extension via settings object with values in extensions",
53+
extensions: map[string]any{
54+
testExtURI: map[string]any{"setting": "val"},
55+
},
56+
serverExts: map[string]any{testExtURI: map[string]any{}},
57+
expectedUri: testExtURI,
58+
expectedVal: true,
59+
},
60+
{
61+
name: "server does not support extension",
62+
extensions: map[string]any{
63+
testExtURI: map[string]any{},
64+
},
65+
serverExts: map[string]any{"other-extension": map[string]any{}},
66+
expectedUri: testExtURI,
67+
expectedVal: false,
68+
},
69+
{
70+
name: "nil value in client extensions",
71+
extensions: map[string]any{
72+
testExtURI: nil,
73+
},
74+
serverExts: nil,
75+
expectedUri: testExtURI,
76+
expectedVal: true,
77+
},
78+
{
79+
name: "server extensions empty",
80+
extensions: map[string]any{
81+
testExtURI: map[string]any{},
82+
},
83+
serverExts: map[string]any{},
84+
expectedUri: testExtURI,
85+
expectedVal: false,
86+
},
87+
}
88+
89+
for _, tc := range tests {
90+
t.Run(tc.name, func(t *testing.T) {
91+
origSupported := SupportedExtensions
92+
t.Cleanup(func() {
93+
SupportedExtensions = origSupported
94+
})
95+
if tc.serverExts != nil {
96+
SupportedExtensions = tc.serverExts
97+
} else {
98+
SupportedExtensions = map[string]any{testExtURI: map[string]any{}}
99+
}
100+
Initialize(nil)
101+
exts := ParseSupportedExtensions(tc.extensions)
102+
_, ok := exts[tc.expectedUri]
103+
if ok != tc.expectedVal {
104+
t.Errorf("ParseSupportedExtensions() value for %s = %v, want %v", tc.expectedUri, ok, tc.expectedVal)
105+
}
106+
})
107+
}
108+
}
109+
110+
func TestServerExtensions(t *testing.T) {
111+
orig := ServerExtensions
112+
t.Cleanup(func() {
113+
ServerExtensions = orig
114+
})
115+
ServerExtensions = nil
116+
if ServerExtensions != nil {
117+
t.Errorf("expected nil when no server extensions registered")
118+
}
119+
120+
ServerExtensions = map[string]any{testExtURI: map[string]any{}}
121+
if ServerExtensions == nil || ServerExtensions[testExtURI] == nil {
122+
t.Errorf("expected testExtURI to be registered in server extensions")
123+
}
124+
}

‎internal/server/mcp/v20260728/method.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,7 @@ func serverDiscoverHandler(ctx context.Context, id jsonrpc.RequestId, body []byt
196196
},
197197
SupportedVersions: mcputil.GetSupportedVersions(enableDraft),
198198
Capabilities: ServerCapabilities{
199+
Extensions: ServerExtensions,
199200
Tools: &ListChanged{
200201
ListChanged: &toolsListChanged,
201202
},

‎internal/server/mcp/v20260728/types.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,9 @@ type ResultMetaObject struct {
118118
// client can define its own, additional capabilities.
119119
type ClientCapabilities struct {
120120
// Experimental, non-standard capabilities that the client supports.
121-
Experimental map[string]interface{} `json:"experimental,omitempty"`
121+
Experimental map[string]any `json:"experimental,omitempty"`
122+
// Standard extensions that the client supports.
123+
Extensions map[string]any `json:"extensions,omitempty"`
122124
// Present if the client supports listing roots.
123125
Roots *ListChanged `json:"roots,omitempty"`
124126
// Present if the client supports sampling from an LLM.
@@ -185,8 +187,9 @@ type Implementation struct {
185187
// capabilities are defined here, in this schema, but this is not a closed set: any
186188
// server can define its own, additional capabilities.
187189
type ServerCapabilities struct {
188-
Tools *ListChanged `json:"tools,omitempty"`
189-
Prompts *ListChanged `json:"prompts,omitempty"`
190+
Extensions map[string]any `json:"extensions,omitempty"`
191+
Tools *ListChanged `json:"tools,omitempty"`
192+
Prompts *ListChanged `json:"prompts,omitempty"`
190193
}
191194

192195
// ListChange represents whether the server supports notification for changes to the capabilities.

‎internal/server/server.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import (
3939
"github.com/googleapis/mcp-toolbox/internal/group"
4040
"github.com/googleapis/mcp-toolbox/internal/log"
4141
"github.com/googleapis/mcp-toolbox/internal/prompts"
42+
"github.com/googleapis/mcp-toolbox/internal/server/mcp"
4243
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
4344
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
4445
"github.com/googleapis/mcp-toolbox/internal/sources"
@@ -469,6 +470,10 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
469470
limit = DefaultHTTPMaxRequestBytes
470471
}
471472

473+
mcp.InitializeProtocols(mcp.ProtocolOptions{
474+
DisableExt: cfg.DisableExt,
475+
})
476+
472477
s := &Server{
473478
version: cfg.Version,
474479
sqlCommenterEnabled: cfg.SQLCommenter,

‎internal/server/server_test.go‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ import (
4747
"github.com/googleapis/mcp-toolbox/internal/prompts"
4848
_ "github.com/googleapis/mcp-toolbox/internal/prompts/custom"
4949
"github.com/googleapis/mcp-toolbox/internal/server"
50+
v20260728 "github.com/googleapis/mcp-toolbox/internal/server/mcp/v20260728"
5051
"github.com/googleapis/mcp-toolbox/internal/sources"
5152
"github.com/googleapis/mcp-toolbox/internal/sources/alloydbpg"
5253
"github.com/googleapis/mcp-toolbox/internal/telemetry"
@@ -1980,3 +1981,72 @@ func TestDefaultToolsetIsAlphabeticallySorted(t *testing.T) {
19801981
t.Errorf("default toolset ToolNames mismatch (-want +got):\n%s", diff)
19811982
}
19821983
}
1984+
1985+
func TestNewServer_Extensions(t *testing.T) {
1986+
orig := v20260728.SupportedExtensions
1987+
t.Cleanup(func() {
1988+
v20260728.SupportedExtensions = orig
1989+
})
1990+
v20260728.SupportedExtensions = map[string]any{"com.google.cloud/toolbox.v1": map[string]any{}, "io.modelcontextprotocol/tasks": map[string]any{}}
1991+
1992+
ctx := context.Background()
1993+
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
1994+
if err != nil {
1995+
t.Fatalf("unexpected error: %s", err)
1996+
}
1997+
ctx = util.WithLogger(ctx, testLogger)
1998+
1999+
instrumentation, err := telemetry.CreateTelemetryInstrumentation("0.0.0")
2000+
if err != nil {
2001+
t.Fatalf("unexpected error: %s", err)
2002+
}
2003+
ctx = util.WithInstrumentation(ctx, instrumentation)
2004+
2005+
tests := []struct {
2006+
name string
2007+
disableExt []string
2008+
want []string
2009+
}{
2010+
{
2011+
name: "default enables all supported extensions",
2012+
disableExt: nil,
2013+
want: []string{"com.google.cloud/toolbox.v1", "io.modelcontextprotocol/tasks"},
2014+
},
2015+
{
2016+
name: "disable one extension",
2017+
disableExt: []string{"io.modelcontextprotocol/tasks"},
2018+
want: []string{"com.google.cloud/toolbox.v1"},
2019+
},
2020+
{
2021+
name: "disable all supported extensions",
2022+
disableExt: []string{"com.google.cloud/toolbox.v1", "io.modelcontextprotocol/tasks"},
2023+
want: nil,
2024+
},
2025+
{
2026+
name: "empty strings or unknown extensions in disableExt ignored",
2027+
disableExt: []string{"", "com.example/unknown"},
2028+
want: []string{"com.google.cloud/toolbox.v1", "io.modelcontextprotocol/tasks"},
2029+
},
2030+
}
2031+
2032+
for _, tt := range tests {
2033+
t.Run(tt.name, func(t *testing.T) {
2034+
cfg := server.ServerConfig{
2035+
DisableExt: tt.disableExt,
2036+
}
2037+
_, err := server.NewServer(ctx, cfg)
2038+
if err != nil {
2039+
t.Fatalf("unexpected error: %v", err)
2040+
}
2041+
var got []string
2042+
for k := range v20260728.ServerExtensions {
2043+
got = append(got, k)
2044+
}
2045+
slices.Sort(got)
2046+
slices.Sort(tt.want)
2047+
if !reflect.DeepEqual(got, tt.want) {
2048+
t.Errorf("v20260728.ServerExtensions keys = %v, want %v", got, tt.want)
2049+
}
2050+
})
2051+
}
2052+
}

0 commit comments

Comments
 (0)