Skip to content

Commit 25ce953

Browse files
Yuan325Yuan325
authored
fix: re-add tool validation during startup (#3705)
Tool validation is added during server startup, except for during skills generation.
1 parent cf128ff commit 25ce953

15 files changed

Lines changed: 219 additions & 179 deletions

File tree

‎cmd/internal/skills/generator_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,7 @@ func TestFormatParameters(t *testing.T) {
175175

176176
func TestGenerateSkillMarkdown(t *testing.T) {
177177
toolsMap := map[string]tools.Tool{
178-
"tool1": testutils.NewMockTool("tool1", "First tool",
178+
"tool1": testutils.NewMockTool("tool1", "First tool", "",
179179
[]parameters.Parameter{
180180
parameters.NewStringParameter("p1", "d1"),
181181
}, false, false),

‎internal/group/group_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ import (
2929

3030
func testFixtures() (map[string]tools.Tool, map[string]prompts.Prompt) {
3131
toolsMap := map[string]tools.Tool{
32-
"tool1": testutils.NewMockTool("tool1", "first tool", []parameters.Parameter{}, false, false),
33-
"tool2": testutils.NewMockTool("tool2", "second tool", []parameters.Parameter{}, false, false),
32+
"tool1": testutils.NewMockTool("tool1", "first tool", "", []parameters.Parameter{}, false, false),
33+
"tool2": testutils.NewMockTool("tool2", "second tool", "", []parameters.Parameter{}, false, false),
3434
}
3535
promptsMap := map[string]prompts.Prompt{
3636
"prompt1": testutils.NewMockPrompt("prompt1", "first prompt", prompts.Arguments{}),

‎internal/server/api_test.go‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ func TestToolsetEndpoint(t *testing.T) {
5454
want: wantResponse{
5555
statusCode: http.StatusOK,
5656
version: testutils.MockVersionString,
57-
tools: []string{testutils.MockTool1.Name, testutils.MockTool2.Name},
57+
tools: []string{testutils.MockTool1.GetName(), testutils.MockTool2.GetName()},
5858
},
5959
},
6060
{
@@ -71,7 +71,7 @@ func TestToolsetEndpoint(t *testing.T) {
7171
want: wantResponse{
7272
statusCode: http.StatusOK,
7373
version: testutils.MockVersionString,
74-
tools: []string{testutils.MockTool1.Name},
74+
tools: []string{testutils.MockTool1.GetName()},
7575
},
7676
},
7777
{
@@ -80,7 +80,7 @@ func TestToolsetEndpoint(t *testing.T) {
8080
want: wantResponse{
8181
statusCode: http.StatusOK,
8282
version: testutils.MockVersionString,
83-
tools: []string{testutils.MockTool2.Name},
83+
tools: []string{testutils.MockTool2.GetName()},
8484
},
8585
},
8686
}
@@ -147,20 +147,20 @@ func TestToolGetEndpoint(t *testing.T) {
147147
}{
148148
{
149149
name: "tool1",
150-
toolName: testutils.MockTool1.Name,
150+
toolName: testutils.MockTool1.GetName(),
151151
want: wantResponse{
152152
statusCode: http.StatusOK,
153153
version: testutils.MockVersionString,
154-
tools: []string{testutils.MockTool1.Name},
154+
tools: []string{testutils.MockTool1.GetName()},
155155
},
156156
},
157157
{
158158
name: "tool2",
159-
toolName: testutils.MockTool2.Name,
159+
toolName: testutils.MockTool2.GetName(),
160160
want: wantResponse{
161161
statusCode: http.StatusOK,
162162
version: testutils.MockVersionString,
163-
tools: []string{testutils.MockTool2.Name},
163+
tools: []string{testutils.MockTool2.GetName()},
164164
},
165165
},
166166
{
@@ -228,15 +228,15 @@ func TestToolInvokeEndpoint(t *testing.T) {
228228
isErr bool
229229
}{
230230
{
231-
name: "tool1",
232-
toolName: testutils.MockTool1.Name,
231+
name: "tool without param",
232+
toolName: testutils.MockTool1.GetName(),
233233
requestBody: bytes.NewBuffer([]byte(`{}`)),
234234
want: "{result:[no_params]}\n",
235235
isErr: false,
236236
},
237237
{
238-
name: "tool2",
239-
toolName: testutils.MockTool2.Name,
238+
name: "tool with params",
239+
toolName: testutils.MockTool2.GetName(),
240240
requestBody: bytes.NewBuffer([]byte(`{"param1": 1, "param2": 2}`)),
241241
want: "{result:[some_params]}\n",
242242
isErr: false,
@@ -249,15 +249,15 @@ func TestToolInvokeEndpoint(t *testing.T) {
249249
isErr: true,
250250
},
251251
{
252-
name: "tool4",
253-
toolName: testutils.MockTool4.Name,
252+
name: "unauthorized tools",
253+
toolName: testutils.MockTool4.GetName(),
254254
requestBody: bytes.NewBuffer([]byte(`{}`)),
255255
want: "",
256256
isErr: true,
257257
},
258258
{
259-
name: "tool5",
260-
toolName: testutils.MockTool5.Name,
259+
name: "tool requiring client auth",
260+
toolName: testutils.MockTool5.GetName(),
261261
requestBody: bytes.NewBuffer([]byte(`{}`)),
262262
want: "",
263263
isErr: true,
@@ -307,7 +307,7 @@ func TestApiRequestBodyLimit(t *testing.T) {
307307

308308
limit := int(DefaultHTTPMaxRequestBytes)
309309
tooLarge := []byte(fmt.Sprintf(`{"param":"%s"}`, strings.Repeat("x", limit)))
310-
resp, body, err := runRequest(ts, http.MethodPost, fmt.Sprintf("/tool/%s/invoke", testutils.MockTool1.Name), bytes.NewReader(tooLarge), nil)
310+
resp, body, err := runRequest(ts, http.MethodPost, fmt.Sprintf("/tool/%s/invoke", testutils.MockTool1.GetName()), bytes.NewReader(tooLarge), nil)
311311
if err != nil {
312312
t.Fatalf("unexpected error during request: %s", err)
313313
}
@@ -339,7 +339,7 @@ func TestApiRequestBodyLimitOverride(t *testing.T) {
339339
defer ts.Close()
340340

341341
tooLarge := []byte(fmt.Sprintf(`{"param":"%s"}`, strings.Repeat("x", int(customLimit))))
342-
resp, body, err := runRequest(ts, http.MethodPost, fmt.Sprintf("/tool/%s/invoke", testutils.MockTool1.Name), bytes.NewReader(tooLarge), nil)
342+
resp, body, err := runRequest(ts, http.MethodPost, fmt.Sprintf("/tool/%s/invoke", testutils.MockTool1.GetName()), bytes.NewReader(tooLarge), nil)
343343
if err != nil {
344344
t.Fatalf("unexpected error during request: %s", err)
345345
}

‎internal/server/config.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,8 @@ type ServerConfig struct {
101101
HttpMaxRequestBytes int64
102102
// EnableDraftSpecs allow users to opt-in and test upcoming draft MCP specs.
103103
EnableDraftSpecs bool
104+
// SkipSourceValidation skips source validation during server startup
105+
SkipSourceValidation bool
104106
}
105107

106108
type logFormat string

‎internal/server/mcp/v20241105/manifests_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -231,17 +231,17 @@ func TestParamManifest(t *testing.T) {
231231
}
232232

233233
func TestGenerateListToolsResult(t *testing.T) {
234-
tool1 := testutils.NewMockTool("no_params", "", []parameters.Parameter{}, false, false)
234+
tool1 := testutils.NewMockTool("no_params", "", "", []parameters.Parameter{}, false, false)
235235
tool2 := testutils.NewMockTool(
236236
"some_params",
237-
"",
237+
"", "",
238238
parameters.Parameters{
239239
parameters.NewIntParameter("param1", "This is the first parameter."),
240240
parameters.NewIntParameter("param2", "This is the second parameter."),
241241
}, false, false)
242242
toolsMap := make(map[string]tools.Tool)
243-
toolsMap[tool1.Name] = tool1
244-
toolsMap[tool2.Name] = tool2
243+
toolsMap[tool1.GetName()] = tool1
244+
toolsMap[tool2.GetName()] = tool2
245245
g := group.NewGroup(group.GroupConfig{
246246
Name: "test-toolset",
247247
ToolNames: []string{"no_params", "some_params"},

‎internal/server/mcp/v20250326/manifests_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -231,17 +231,17 @@ func TestParamManifest(t *testing.T) {
231231
}
232232

233233
func TestGenerateListToolsResult(t *testing.T) {
234-
tool1 := testutils.NewMockTool("no_params", "", []parameters.Parameter{}, false, false)
234+
tool1 := testutils.NewMockTool("no_params", "", "", []parameters.Parameter{}, false, false)
235235
tool2 := testutils.NewMockTool(
236236
"some_params",
237-
"",
237+
"", "",
238238
parameters.Parameters{
239239
parameters.NewIntParameter("param1", "This is the first parameter."),
240240
parameters.NewIntParameter("param2", "This is the second parameter."),
241241
}, false, false)
242242
toolsMap := make(map[string]tools.Tool)
243-
toolsMap[tool1.Name] = tool1
244-
toolsMap[tool2.Name] = tool2
243+
toolsMap[tool1.GetName()] = tool1
244+
toolsMap[tool2.GetName()] = tool2
245245
g := group.NewGroup(group.GroupConfig{
246246
Name: "test-toolset",
247247
ToolNames: []string{"no_params", "some_params"},

‎internal/server/mcp/v20250618/manifests_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -231,17 +231,17 @@ func TestParamManifest(t *testing.T) {
231231
}
232232

233233
func TestGenerateListToolsResult(t *testing.T) {
234-
tool1 := testutils.NewMockTool("no_params", "", []parameters.Parameter{}, false, false)
234+
tool1 := testutils.NewMockTool("no_params", "", "", []parameters.Parameter{}, false, false)
235235
tool2 := testutils.NewMockTool(
236236
"some_params",
237-
"",
237+
"", "",
238238
parameters.Parameters{
239239
parameters.NewIntParameter("param1", "This is the first parameter."),
240240
parameters.NewIntParameter("param2", "This is the second parameter."),
241241
}, false, false)
242242
toolsMap := make(map[string]tools.Tool)
243-
toolsMap[tool1.Name] = tool1
244-
toolsMap[tool2.Name] = tool2
243+
toolsMap[tool1.GetName()] = tool1
244+
toolsMap[tool2.GetName()] = tool2
245245
g := group.NewGroup(group.GroupConfig{
246246
Name: "test-toolset",
247247
ToolNames: []string{"no_params", "some_params"},

‎internal/server/mcp/v20251125/manifests_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -231,17 +231,17 @@ func TestParamManifest(t *testing.T) {
231231
}
232232

233233
func TestGenerateListToolsResult(t *testing.T) {
234-
tool1 := testutils.NewMockTool("no_params", "", []parameters.Parameter{}, false, false)
234+
tool1 := testutils.NewMockTool("no_params", "", "", []parameters.Parameter{}, false, false)
235235
tool2 := testutils.NewMockTool(
236236
"some_params",
237-
"",
237+
"", "",
238238
parameters.Parameters{
239239
parameters.NewIntParameter("param1", "This is the first parameter."),
240240
parameters.NewIntParameter("param2", "This is the second parameter."),
241241
}, false, false)
242242
toolsMap := make(map[string]tools.Tool)
243-
toolsMap[tool1.Name] = tool1
244-
toolsMap[tool2.Name] = tool2
243+
toolsMap[tool1.GetName()] = tool1
244+
toolsMap[tool2.GetName()] = tool2
245245
g := group.NewGroup(group.GroupConfig{
246246
Name: "test-toolset",
247247
ToolNames: []string{"no_params", "some_params"},

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -231,17 +231,17 @@ func TestParamManifest(t *testing.T) {
231231
}
232232

233233
func TestGenerateListToolsResult(t *testing.T) {
234-
tool1 := testutils.NewMockTool("no_params", "", []parameters.Parameter{}, false, false)
234+
tool1 := testutils.NewMockTool("no_params", "", "", []parameters.Parameter{}, false, false)
235235
tool2 := testutils.NewMockTool(
236236
"some_params",
237-
"",
237+
"", "",
238238
parameters.Parameters{
239239
parameters.NewIntParameter("param1", "This is the first parameter."),
240240
parameters.NewIntParameter("param2", "This is the second parameter."),
241241
}, false, false)
242242
toolsMap := make(map[string]tools.Tool)
243-
toolsMap[tool1.Name] = tool1
244-
toolsMap[tool2.Name] = tool2
243+
toolsMap[tool1.GetName()] = tool1
244+
toolsMap[tool2.GetName()] = tool2
245245
g := group.NewGroup(group.GroupConfig{
246246
Name: "test-toolset",
247247
ToolNames: []string{"no_params", "some_params"},

‎internal/server/server.go‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
183183
}
184184
l.InfoContext(ctx, fmt.Sprintf("Initialized %d embeddingModels: %s", len(embeddingModelsMap), strings.Join(embeddingModelNames, ", ")))
185185

186-
toolsMap, err := initializeTools(ctx, cfg, instrumentation, l)
186+
toolsMap, err := initializeTools(ctx, cfg, sourcesMap, instrumentation, l)
187187
if err != nil {
188188
return nil, nil, nil, nil, nil, nil, err
189189
}
@@ -241,8 +241,9 @@ func InitializeOfflineConfigs(ctx context.Context, cfg ServerConfig) (
241241
if err != nil {
242242
return nil, nil, fmt.Errorf("failed to get logger from context: %w", err)
243243
}
244-
245-
toolsMap, err := initializeTools(ctx, cfg, instrumentation, l)
244+
// Automatically skip source validation
245+
cfg.SkipSourceValidation = true
246+
toolsMap, err := initializeTools(ctx, cfg, map[string]sources.Source{}, instrumentation, l)
246247
if err != nil {
247248
return nil, nil, err
248249
}
@@ -266,7 +267,7 @@ func InitializeOfflineConfigs(ctx context.Context, cfg ServerConfig) (
266267
}
267268

268269
// initializeTools initializes and validates the tools from the config.
269-
func initializeTools(ctx context.Context, cfg ServerConfig, instrumentation *telemetry.Instrumentation, l log.Logger) (map[string]tools.Tool, error) {
270+
func initializeTools(ctx context.Context, cfg ServerConfig, sourcesMap map[string]sources.Source, instrumentation *telemetry.Instrumentation, l log.Logger) (map[string]tools.Tool, error) {
270271
toolsMap := make(map[string]tools.Tool)
271272
for name, tc := range cfg.ToolConfigs {
272273
t, err := func() (tools.Tool, error) {
@@ -281,6 +282,21 @@ func initializeTools(ctx context.Context, cfg ServerConfig, instrumentation *tel
281282
if err != nil {
282283
return nil, fmt.Errorf("unable to initialize tool %q: %w", name, err)
283284
}
285+
if !cfg.SkipSourceValidation {
286+
srcName := t.GetSourceName()
287+
var src sources.Source
288+
var ok bool
289+
if srcName != "" {
290+
src, ok = sourcesMap[srcName]
291+
if !ok {
292+
return nil, fmt.Errorf("unable to retrieve source %s for tool %s", srcName, name)
293+
}
294+
}
295+
err = t.ValidateSource(src)
296+
if err != nil {
297+
return nil, err
298+
}
299+
}
284300
return t, nil
285301
}()
286302
if err != nil {

0 commit comments

Comments
 (0)