Skip to content

Commit 209743a

Browse files
ChrisJBurnsChrisJBurns
committed
uses schema
Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com>
1 parent b2e2680 commit 209743a

File tree

3 files changed

+89
-83
lines changed

3 files changed

+89
-83
lines changed

‎pkg/authz/response_filter.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -269,7 +269,8 @@ func (rfw *ResponseFilteringWriter) filterToolsResponse(response *jsonrpc2.Respo
269269

270270
// Create a new result with filtered tools
271271
filteredResult := mcp.ListToolsResult{
272-
Tools: filteredTools,
272+
NextCursor: listResult.NextCursor,
273+
Tools: filteredTools,
273274
}
274275

275276
// Marshal the filtered result back
@@ -320,7 +321,8 @@ func (rfw *ResponseFilteringWriter) filterPromptsResponse(response *jsonrpc2.Res
320321

321322
// Create a new result with filtered prompts
322323
filteredResult := mcp.ListPromptsResult{
323-
Prompts: filteredPrompts,
324+
NextCursor: listResult.NextCursor,
325+
Prompts: filteredPrompts,
324326
}
325327

326328
// Marshal the filtered result back
@@ -371,7 +373,8 @@ func (rfw *ResponseFilteringWriter) filterResourcesResponse(response *jsonrpc2.R
371373

372374
// Create a new result with filtered resources
373375
filteredResult := mcp.ListResourcesResult{
374-
Resources: filteredResources,
376+
NextCursor: listResult.NextCursor,
377+
Resources: filteredResources,
375378
}
376379

377380
// Marshal the filtered result back

‎pkg/mcp/parser_integration_test.go‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"testing"
99
"time"
1010

11+
"github.com/google/jsonschema-go/jsonschema"
1112
"github.com/modelcontextprotocol/go-sdk/mcp"
1213
"github.com/stretchr/testify/assert"
1314
"github.com/stretchr/testify/require"
@@ -316,15 +317,16 @@ func createTestMCPServer() *mcp.Server {
316317
&mcp.Tool{
317318
Name: "test_tool",
318319
Description: "A test tool",
319-
InputSchema: json.RawMessage(`{
320-
"type": "object",
321-
"properties": {
320+
InputSchema: &jsonschema.Schema{
321+
Type: "object",
322+
Properties: map[string]*jsonschema.Schema{
322323
"message": {
323-
"type": "string",
324-
"description": "Test message"
325-
}
326-
}
327-
}`),
324+
Type: "string",
325+
Description: "Test message",
326+
},
327+
},
328+
Required: []string{"message"},
329+
},
328330
},
329331
func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
330332
return &mcp.CallToolResult{

‎pkg/mcp/server/server.go‎

Lines changed: 73 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"net/http"
77
"time"
88

9+
"github.com/google/jsonschema-go/jsonschema"
910
"github.com/modelcontextprotocol/go-sdk/mcp"
1011

1112
"github.com/stacklok/toolhive/pkg/logger"
@@ -104,141 +105,141 @@ func registerTools(mcpServer *mcp.Server, handler *Handler) {
104105
mcpServer.AddTool(&mcp.Tool{
105106
Name: "search_registry",
106107
Description: "Search the ToolHive registry for MCP servers",
107-
InputSchema: map[string]interface{}{
108-
"type": "object",
109-
"properties": map[string]interface{}{
110-
"query": map[string]interface{}{
111-
"type": "string",
112-
"description": "Search query to find MCP servers",
108+
InputSchema: &jsonschema.Schema{
109+
Type: "object",
110+
Properties: map[string]*jsonschema.Schema{
111+
"query": {
112+
Type: "string",
113+
Description: "Search query to find MCP servers",
113114
},
114115
},
115-
"required": []string{"query"},
116+
Required: []string{"query"},
116117
},
117118
}, handler.SearchRegistry)
118119

119120
mcpServer.AddTool(&mcp.Tool{
120121
Name: "run_server",
121122
Description: "Run an MCP server from the ToolHive registry",
122-
InputSchema: map[string]interface{}{
123-
"type": "object",
124-
"properties": map[string]interface{}{
125-
"server": map[string]interface{}{
126-
"type": "string",
127-
"description": "Name of the server to run (e.g., 'fetch', 'github')",
123+
InputSchema: &jsonschema.Schema{
124+
Type: "object",
125+
Properties: map[string]*jsonschema.Schema{
126+
"server": {
127+
Type: "string",
128+
Description: "Name of the server to run (e.g., 'fetch', 'github')",
128129
},
129-
"name": map[string]interface{}{
130-
"type": "string",
131-
"description": "Optional custom name for the server instance",
130+
"name": {
131+
Type: "string",
132+
Description: "Optional custom name for the server instance",
132133
},
133-
"env": map[string]interface{}{
134-
"type": "object",
135-
"description": "Environment variables to pass to the server",
136-
"additionalProperties": map[string]interface{}{
137-
"type": "string",
134+
"env": {
135+
Type: "object",
136+
Description: "Environment variables to pass to the server",
137+
AdditionalProperties: &jsonschema.Schema{
138+
Type: "string",
138139
},
139140
},
140-
"secrets": map[string]interface{}{
141-
"type": "array",
142-
"description": "Secrets to pass to the server as environment variables",
143-
"items": map[string]interface{}{
144-
"type": "object",
145-
"properties": map[string]interface{}{
146-
"name": map[string]interface{}{
147-
"type": "string",
148-
"description": "Name of the secret in the ToolHive secrets store",
141+
"secrets": {
142+
Type: "array",
143+
Description: "Secrets to pass to the server as environment variables",
144+
Items: &jsonschema.Schema{
145+
Type: "object",
146+
Properties: map[string]*jsonschema.Schema{
147+
"name": {
148+
Type: "string",
149+
Description: "Name of the secret in the ToolHive secrets store",
149150
},
150-
"target": map[string]interface{}{
151-
"type": "string",
152-
"description": "Target environment variable name in the server container",
151+
"target": {
152+
Type: "string",
153+
Description: "Target environment variable name in the server container",
153154
},
154155
},
155-
"required": []string{"name", "target"},
156+
Required: []string{"name", "target"},
156157
},
157158
},
158159
},
159-
"required": []string{"server"},
160+
Required: []string{"server"},
160161
},
161162
}, handler.RunServer)
162163

163164
mcpServer.AddTool(&mcp.Tool{
164165
Name: "list_servers",
165166
Description: "List all running ToolHive MCP servers",
166-
InputSchema: map[string]interface{}{
167-
"type": "object",
168-
"properties": map[string]interface{}{},
167+
InputSchema: &jsonschema.Schema{
168+
Type: "object",
169+
Properties: map[string]*jsonschema.Schema{},
169170
},
170171
}, handler.ListServers)
171172

172173
mcpServer.AddTool(&mcp.Tool{
173174
Name: "stop_server",
174175
Description: "Stop a running MCP server",
175-
InputSchema: map[string]interface{}{
176-
"type": "object",
177-
"properties": map[string]interface{}{
178-
"name": map[string]interface{}{
179-
"type": "string",
180-
"description": "Name of the server to stop",
176+
InputSchema: &jsonschema.Schema{
177+
Type: "object",
178+
Properties: map[string]*jsonschema.Schema{
179+
"name": {
180+
Type: "string",
181+
Description: "Name of the server to stop",
181182
},
182183
},
183-
"required": []string{"name"},
184+
Required: []string{"name"},
184185
},
185186
}, handler.StopServer)
186187

187188
mcpServer.AddTool(&mcp.Tool{
188189
Name: "remove_server",
189190
Description: "Remove a stopped MCP server",
190-
InputSchema: map[string]interface{}{
191-
"type": "object",
192-
"properties": map[string]interface{}{
193-
"name": map[string]interface{}{
194-
"type": "string",
195-
"description": "Name of the server to remove",
191+
InputSchema: &jsonschema.Schema{
192+
Type: "object",
193+
Properties: map[string]*jsonschema.Schema{
194+
"name": {
195+
Type: "string",
196+
Description: "Name of the server to remove",
196197
},
197198
},
198-
"required": []string{"name"},
199+
Required: []string{"name"},
199200
},
200201
}, handler.RemoveServer)
201202

202203
mcpServer.AddTool(&mcp.Tool{
203204
Name: "get_server_logs",
204205
Description: "Get logs from a running MCP server",
205-
InputSchema: map[string]interface{}{
206-
"type": "object",
207-
"properties": map[string]interface{}{
208-
"name": map[string]interface{}{
209-
"type": "string",
210-
"description": "Name of the server to get logs from",
206+
InputSchema: &jsonschema.Schema{
207+
Type: "object",
208+
Properties: map[string]*jsonschema.Schema{
209+
"name": {
210+
Type: "string",
211+
Description: "Name of the server to get logs from",
211212
},
212213
},
213-
"required": []string{"name"},
214+
Required: []string{"name"},
214215
},
215216
}, handler.GetServerLogs)
216217

217218
mcpServer.AddTool(&mcp.Tool{
218219
Name: "list_secrets",
219220
Description: "List all available secrets in the ToolHive secrets store",
220-
InputSchema: map[string]interface{}{
221-
"type": "object",
222-
"properties": map[string]interface{}{},
221+
InputSchema: &jsonschema.Schema{
222+
Type: "object",
223+
Properties: map[string]*jsonschema.Schema{},
223224
},
224225
}, handler.ListSecrets)
225226

226227
mcpServer.AddTool(&mcp.Tool{
227228
Name: "set_secret",
228229
Description: "Set a secret by reading its value from a file",
229-
InputSchema: map[string]interface{}{
230-
"type": "object",
231-
"properties": map[string]interface{}{
232-
"name": map[string]interface{}{
233-
"type": "string",
234-
"description": "Name of the secret to set",
230+
InputSchema: &jsonschema.Schema{
231+
Type: "object",
232+
Properties: map[string]*jsonschema.Schema{
233+
"name": {
234+
Type: "string",
235+
Description: "Name of the secret to set",
235236
},
236-
"file_path": map[string]interface{}{
237-
"type": "string",
238-
"description": "Path to the file containing the secret value",
237+
"file_path": {
238+
Type: "string",
239+
Description: "Path to the file containing the secret value",
239240
},
240241
},
241-
"required": []string{"name", "file_path"},
242+
Required: []string{"name", "file_path"},
242243
},
243244
}, handler.SetSecret)
244245
}

0 commit comments

Comments
 (0)