Skip to content

Commit 0e4effa

Browse files
sinadarbouysinadarbouy
committed
Add optional flags to stream commands in StreamTools
- Introduced a new "flags" parameter in the GetTools method for various stream commands, allowing users to pass optional flags as an array of strings. - Implemented a helper function, getFlags, to extract and handle these flags from command arguments. - Updated stream command handlers to incorporate the new flags functionality, enhancing command flexibility and usability.
1 parent ab39520 commit 0e4effa

File tree

1 file changed

+102
-8
lines changed

1 file changed

+102
-8
lines changed

‎tools/stream.go

Lines changed: 102 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,11 @@ func (s *StreamTools) GetTools() []Tool {
3939
"type": "string",
4040
"description": "The name of the stream to get information about",
4141
},
42+
"flags": map[string]interface{}{
43+
"type": "array",
44+
"items": map[string]interface{}{"type": "string"},
45+
"description": "Optional flags to pass to the command",
46+
},
4247
},
4348
Required: []string{"account_name", "stream"},
4449
},
@@ -56,6 +61,11 @@ func (s *StreamTools) GetTools() []Tool {
5661
"type": "string",
5762
"description": "The NATS account to use",
5863
},
64+
"flags": map[string]interface{}{
65+
"type": "array",
66+
"items": map[string]interface{}{"type": "string"},
67+
"description": "Optional flags to pass to the command",
68+
},
5969
},
6070
Required: []string{"account_name"},
6171
},
@@ -73,6 +83,11 @@ func (s *StreamTools) GetTools() []Tool {
7383
"type": "string",
7484
"description": "The NATS account to use",
7585
},
86+
"flags": map[string]interface{}{
87+
"type": "array",
88+
"items": map[string]interface{}{"type": "string"},
89+
"description": "Optional flags to pass to the command",
90+
},
7691
},
7792
Required: []string{"account_name"},
7893
},
@@ -90,6 +105,11 @@ func (s *StreamTools) GetTools() []Tool {
90105
"type": "string",
91106
"description": "The NATS account to use",
92107
},
108+
"flags": map[string]interface{}{
109+
"type": "array",
110+
"items": map[string]interface{}{"type": "string"},
111+
"description": "Optional flags to pass to the command",
112+
},
93113
},
94114
Required: []string{"account_name"},
95115
},
@@ -111,6 +131,11 @@ func (s *StreamTools) GetTools() []Tool {
111131
"type": "string",
112132
"description": "The name of the stream to get state for",
113133
},
134+
"flags": map[string]interface{}{
135+
"type": "array",
136+
"items": map[string]interface{}{"type": "string"},
137+
"description": "Optional flags to pass to the command",
138+
},
114139
},
115140
Required: []string{"account_name", "stream"},
116141
},
@@ -132,6 +157,11 @@ func (s *StreamTools) GetTools() []Tool {
132157
"type": "string",
133158
"description": "Stream name",
134159
},
160+
"flags": map[string]interface{}{
161+
"type": "array",
162+
"items": map[string]interface{}{"type": "string"},
163+
"description": "Optional flags to pass to the command",
164+
},
135165
},
136166
Required: []string{"account_name", "stream"},
137167
},
@@ -157,6 +187,11 @@ func (s *StreamTools) GetTools() []Tool {
157187
"type": "integer",
158188
"description": "Page size",
159189
},
190+
"flags": map[string]interface{}{
191+
"type": "array",
192+
"items": map[string]interface{}{"type": "string"},
193+
"description": "Optional flags to pass to the command",
194+
},
160195
},
161196
Required: []string{"account_name", "stream", "size"},
162197
},
@@ -182,6 +217,11 @@ func (s *StreamTools) GetTools() []Tool {
182217
"type": "string",
183218
"description": "Message Sequence to retrieve",
184219
},
220+
"flags": map[string]interface{}{
221+
"type": "array",
222+
"items": map[string]interface{}{"type": "string"},
223+
"description": "Optional flags to pass to the command",
224+
},
185225
},
186226
Required: []string{"account_name", "stream", "id"},
187227
},
@@ -191,6 +231,20 @@ func (s *StreamTools) GetTools() []Tool {
191231
}
192232
}
193233

234+
// Helper function to get flags from arguments
235+
func getFlags(args map[string]interface{}) []string {
236+
if flags, ok := args["flags"].([]interface{}); ok {
237+
strFlags := make([]string, len(flags))
238+
for i, flag := range flags {
239+
if strFlag, ok := flag.(string); ok {
240+
strFlags[i] = strFlag
241+
}
242+
}
243+
return strFlags
244+
}
245+
return nil
246+
}
247+
194248
func (s *StreamTools) streamInfoHandler() server.ToolHandlerFunc {
195249
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
196250
accountName, ok := request.Params.Arguments["account_name"].(string)
@@ -208,7 +262,12 @@ func (s *StreamTools) streamInfoHandler() server.ToolHandlerFunc {
208262
return nil, err
209263
}
210264

211-
output, err := executor.ExecuteCommand("stream", "info", stream)
265+
args := []string{"stream", "info", stream}
266+
if flags := getFlags(request.Params.Arguments); flags != nil {
267+
args = append(args, flags...)
268+
}
269+
270+
output, err := executor.ExecuteCommand(args...)
212271
if err != nil {
213272
return nil, err
214273
}
@@ -228,7 +287,12 @@ func (s *StreamTools) streamListHandler() server.ToolHandlerFunc {
228287
return nil, err
229288
}
230289

231-
output, err := executor.ExecuteCommand("stream", "list")
290+
args := []string{"stream", "list"}
291+
if flags := getFlags(request.Params.Arguments); flags != nil {
292+
args = append(args, flags...)
293+
}
294+
295+
output, err := executor.ExecuteCommand(args...)
232296
if err != nil {
233297
return nil, err
234298
}
@@ -249,7 +313,12 @@ func (s *StreamTools) streamReportHandler() server.ToolHandlerFunc {
249313
return nil, err
250314
}
251315

252-
output, err := executor.ExecuteCommand("stream", "report")
316+
args := []string{"stream", "report"}
317+
if flags := getFlags(request.Params.Arguments); flags != nil {
318+
args = append(args, flags...)
319+
}
320+
321+
output, err := executor.ExecuteCommand(args...)
253322
if err != nil {
254323
return nil, err
255324
}
@@ -270,7 +339,12 @@ func (s *StreamTools) streamFindHandler() server.ToolHandlerFunc {
270339
return nil, err
271340
}
272341

273-
output, err := executor.ExecuteCommand("stream", "find")
342+
args := []string{"stream", "find"}
343+
if flags := getFlags(request.Params.Arguments); flags != nil {
344+
args = append(args, flags...)
345+
}
346+
347+
output, err := executor.ExecuteCommand(args...)
274348
if err != nil {
275349
return nil, err
276350
}
@@ -299,7 +373,12 @@ func (s *StreamTools) streamStateHandler() server.ToolHandlerFunc {
299373
return nil, err
300374
}
301375

302-
output, err := executor.ExecuteCommand("stream", "state", stream)
376+
args := []string{"stream", "state", stream}
377+
if flags := getFlags(request.Params.Arguments); flags != nil {
378+
args = append(args, flags...)
379+
}
380+
381+
output, err := executor.ExecuteCommand(args...)
303382
if err != nil {
304383
return nil, err
305384
}
@@ -328,7 +407,12 @@ func (s *StreamTools) streamSubjectsHandler() server.ToolHandlerFunc {
328407
return nil, err
329408
}
330409

331-
output, err := executor.ExecuteCommand("stream", "subjects", stream)
410+
args := []string{"stream", "subjects", stream}
411+
if flags := getFlags(request.Params.Arguments); flags != nil {
412+
args = append(args, flags...)
413+
}
414+
415+
output, err := executor.ExecuteCommand(args...)
332416
if err != nil {
333417
return nil, err
334418
}
@@ -363,7 +447,12 @@ func (s *StreamTools) streamViewHandler() server.ToolHandlerFunc {
363447
return nil, err
364448
}
365449

366-
output, err := executor.ExecuteCommand("stream", "view", stream, strconv.Itoa(size))
450+
args := []string{"stream", "view", stream, strconv.Itoa(size)}
451+
if flags := getFlags(request.Params.Arguments); flags != nil {
452+
args = append(args, flags...)
453+
}
454+
455+
output, err := executor.ExecuteCommand(args...)
367456
if err != nil {
368457
return nil, err
369458
}
@@ -398,7 +487,12 @@ func (s *StreamTools) streamGetHandler() server.ToolHandlerFunc {
398487
return nil, err
399488
}
400489

401-
output, err := executor.ExecuteCommand("stream", "get", stream, id)
490+
args := []string{"stream", "get", stream, id}
491+
if flags := getFlags(request.Params.Arguments); flags != nil {
492+
args = append(args, flags...)
493+
}
494+
495+
output, err := executor.ExecuteCommand(args...)
402496
if err != nil {
403497
return nil, err
404498
}

0 commit comments

Comments
 (0)