Skip to content

Commit 929714d

Browse files
sinadarbouysinadarbouy
committed
Enhance NATS server tools and configuration
- Refactored NATS server tool handling by introducing separate executors for account and system credentials. - Updated the `newServer` function to accept both account and system NATS credentials. - Expanded the `NATSServerTools` structure to include methods for managing server and stream tools. - Added new stream-related tools for listing, reporting, and retrieving stream information. - Improved command execution handling for NATS CLI commands.
1 parent 228eec9 commit 929714d

File tree

6 files changed

+419
-83
lines changed

6 files changed

+419
-83
lines changed

‎cmd/mcp-nats/main.go

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,10 @@ import (
88
"os"
99

1010
"github.com/mark3labs/mcp-go/server"
11-
"github.com/nats-io/nats.go"
1211
"github.com/sinadarbouy/mcp-nats/tools"
1312
)
1413

15-
func newServer(natsURL string, natsCredsPath string) *server.MCPServer {
14+
func newServer(natsURL string, AccNatsCredsPath string, SysNatsCredsPath string) *server.MCPServer {
1615
s := server.NewMCPServer(
1716
"mcp-nats",
1817
"0.1.0",
@@ -22,16 +21,16 @@ func newServer(natsURL string, natsCredsPath string) *server.MCPServer {
2221
)
2322

2423
// Initialize NATS server tools
25-
natsTools := tools.NewNATSServerTools(natsURL, natsCredsPath)
24+
natsTools := tools.NewNATSServerTools(natsURL, AccNatsCredsPath, SysNatsCredsPath)
2625

2726
// Register all NATS server tools
2827
tools.RegisterTools(s, natsTools)
2928

3029
return s
3130
}
3231

33-
func run(transport, addr, natsURL, natsCredsPath string) error {
34-
s := newServer(natsURL, natsCredsPath)
32+
func run(transport, addr, natsURL, AccNatsCredsPath, SysNatsCredsPath string) error {
33+
s := newServer(natsURL, AccNatsCredsPath, SysNatsCredsPath)
3534

3635
switch transport {
3736
case "stdio":
@@ -57,20 +56,17 @@ func main() {
5756
if natsURL == "" {
5857
log.Fatal("NATS_URL environment variable is required")
5958
}
60-
natsCredsPath := os.Getenv("NATS_CREDS_PATH")
61-
if natsCredsPath == "" {
59+
AccNatsCredsPath := os.Getenv("NATS_CREDS_PATH")
60+
if AccNatsCredsPath == "" {
6261
log.Fatal("NATS_CREDS_PATH environment variable is required")
6362
}
64-
65-
// Connect to NATS
66-
nc, err := nats.Connect(natsURL, nats.UserCredentials(natsCredsPath))
67-
if err != nil {
68-
log.Fatalf("Failed to connect to NATS: %v", err)
63+
SysNatsCredsPath := os.Getenv("NATS_CREDS_PATH_SYS")
64+
if SysNatsCredsPath == "" {
65+
log.Fatal("NATS_CREDS_PATH_SYS environment variable is required")
6966
}
70-
defer nc.Close()
7167

7268
fmt.Printf("Starting MCP NATS server (stdio)...\n")
73-
if err := run("stdio", "0.0.0.0:8002", natsURL, natsCredsPath); err != nil {
69+
if err := run("stdio", "0.0.0.0:8002", natsURL, AccNatsCredsPath, SysNatsCredsPath); err != nil {
7470
panic(err)
7571
}
7672
}

‎tools/common/nats.go

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package common
2+
3+
import (
4+
"fmt"
5+
"os/exec"
6+
)
7+
8+
// NATSExecutor provides common NATS command execution functionality
9+
type NATSExecutor struct {
10+
URL string
11+
CredsPath string
12+
}
13+
14+
// ExecuteCommand executes a NATS CLI command with the configured credentials
15+
func (e *NATSExecutor) ExecuteCommand(args ...string) (string, error) {
16+
baseArgs := []string{"-s", e.URL, "--creds", e.CredsPath}
17+
args = append(baseArgs, args...)
18+
19+
cmd := exec.Command("nats", args...)
20+
output, err := cmd.CombinedOutput()
21+
if err != nil {
22+
return "", fmt.Errorf("NATS command failed: %v, output: %s", err, string(output))
23+
}
24+
return string(output), nil
25+
}

‎tools/nats.go

Lines changed: 45 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,54 @@
11
package tools
22

3+
import "github.com/sinadarbouy/mcp-nats/tools/common"
4+
35
// NATSServerTools contains all NATS server-related tool definitions
46
type NATSServerTools struct {
5-
natsURL string
6-
natsCredsPath string
7+
accExecutor *common.NATSExecutor
8+
sysExecutor *common.NATSExecutor
9+
serverTools *ServerTools
10+
streamTools *StreamTools
711
}
812

913
// NewNATSServerTools creates a new instance of NATSServerTools
10-
func NewNATSServerTools(natsURL, natsCredsPath string) *NATSServerTools {
11-
return &NATSServerTools{
12-
natsURL: natsURL,
13-
natsCredsPath: natsCredsPath,
14+
func NewNATSServerTools(natsURL, AccNatsCredsPath, SysNatsCredsPath string) *NATSServerTools {
15+
accExecutor := &common.NATSExecutor{
16+
URL: natsURL,
17+
CredsPath: AccNatsCredsPath,
18+
}
19+
sysExecutor := &common.NATSExecutor{
20+
URL: natsURL,
21+
CredsPath: SysNatsCredsPath,
22+
}
23+
24+
n := &NATSServerTools{
25+
accExecutor: accExecutor,
26+
sysExecutor: sysExecutor,
1427
}
28+
29+
// Initialize tool categories
30+
n.serverTools = NewServerTools(n)
31+
n.streamTools = NewStreamTools(n)
32+
33+
return n
34+
}
35+
36+
// ServerTools returns the server tools category
37+
func (n *NATSServerTools) ServerTools() ToolCategory {
38+
return n.serverTools
39+
}
40+
41+
// StreamTools returns the stream tools category
42+
func (n *NATSServerTools) StreamTools() ToolCategory {
43+
return n.streamTools
44+
}
45+
46+
// GetSysExecutor returns the system executor
47+
func (n *NATSServerTools) GetSysExecutor() *common.NATSExecutor {
48+
return n.sysExecutor
49+
}
50+
51+
// GetAccExecutor returns the account executor
52+
func (n *NATSServerTools) GetAccExecutor() *common.NATSExecutor {
53+
return n.accExecutor
1554
}

‎tools/server.go

Lines changed: 67 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,26 @@ package tools
22

33
import (
44
"context"
5-
"fmt"
6-
"os/exec"
5+
"strconv"
76

87
"github.com/mark3labs/mcp-go/mcp"
98
"github.com/mark3labs/mcp-go/server"
109
)
1110

12-
// GetServerTools returns all NATS server tools
13-
func (n *NATSServerTools) GetServerTools() []Tool {
11+
// ServerTools represents all NATS server-related tools
12+
type ServerTools struct {
13+
nats *NATSServerTools
14+
}
15+
16+
// NewServerTools creates a new ServerTools instance
17+
func NewServerTools(nats *NATSServerTools) *ServerTools {
18+
return &ServerTools{
19+
nats: nats,
20+
}
21+
}
22+
23+
// GetTools implements the ToolCategory interface
24+
func (s *ServerTools) GetTools() []Tool {
1425
return []Tool{
1526
{
1627
Tool: mcp.Tool{
@@ -19,32 +30,32 @@ func (n *NATSServerTools) GetServerTools() []Tool {
1930
InputSchema: mcp.ToolInputSchema{
2031
Type: "object",
2132
Properties: map[string]interface{}{
22-
"random_string": map[string]interface{}{
23-
"type": "string",
24-
"description": "Dummy parameter for no-parameter tools",
33+
"expect": map[string]interface{}{
34+
"type": "integer",
35+
"description": "How many servers to expect",
2536
},
2637
},
27-
Required: []string{"random_string"},
38+
Required: []string{},
2839
},
2940
},
30-
Handler: n.serverListHandler(),
41+
Handler: s.serverListHandler(),
3142
},
3243
{
3344
Tool: mcp.Tool{
34-
Name: "server_check",
35-
Description: "Check NATS server health",
45+
Name: "server_info",
46+
Description: "Get NATS server info",
3647
InputSchema: mcp.ToolInputSchema{
3748
Type: "object",
3849
Properties: map[string]interface{}{
39-
"random_string": map[string]interface{}{
50+
"server": map[string]interface{}{
4051
"type": "string",
41-
"description": "Dummy parameter for no-parameter tools",
52+
"description": "Server ID or Name to inspect",
4253
},
4354
},
44-
Required: []string{"random_string"},
55+
Required: []string{},
4556
},
4657
},
47-
Handler: n.serverCheckHandler(),
58+
Handler: s.serverInfoHandler(),
4859
},
4960
{
5061
Tool: mcp.Tool{
@@ -53,54 +64,73 @@ func (n *NATSServerTools) GetServerTools() []Tool {
5364
InputSchema: mcp.ToolInputSchema{
5465
Type: "object",
5566
Properties: map[string]interface{}{
56-
"random_string": map[string]interface{}{
57-
"type": "string",
58-
"description": "Dummy parameter for no-parameter tools",
67+
"expect": map[string]interface{}{
68+
"type": "integer",
69+
"description": "How many servers to expect",
5970
},
6071
},
61-
Required: []string{"random_string"},
72+
Required: []string{},
6273
},
6374
},
64-
Handler: n.serverPingHandler(),
75+
Handler: s.serverPingHandler(),
6576
},
6677
}
6778
}
6879

69-
func (n *NATSServerTools) executeNATSCommand(args ...string) (string, error) {
70-
baseArgs := []string{"-s", n.natsURL, "--creds", n.natsCredsPath}
71-
args = append(baseArgs, args...)
80+
// nats server list
81+
// Args:
82+
//
83+
// [<expect>] How many servers to expect
84+
func (s *ServerTools) serverListHandler() server.ToolHandlerFunc {
85+
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
86+
var args []string
87+
args = append(args, "server", "list")
7288

73-
cmd := exec.Command("nats", args...)
74-
output, err := cmd.CombinedOutput()
75-
if err != nil {
76-
return "", fmt.Errorf("NATS command failed: %v, output: %s", err, string(output))
77-
}
78-
return string(output), nil
79-
}
89+
if expect, ok := request.Params.Arguments["expect"].(int); ok {
90+
args = append(args, strconv.Itoa(expect))
91+
}
8092

81-
func (n *NATSServerTools) serverListHandler() server.ToolHandlerFunc {
82-
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
83-
output, err := n.executeNATSCommand("server", "list")
93+
output, err := s.nats.GetSysExecutor().ExecuteCommand(args...)
8494
if err != nil {
8595
return nil, err
8696
}
8797
return mcp.NewToolResultText(output), nil
8898
}
8999
}
90100

91-
func (n *NATSServerTools) serverCheckHandler() server.ToolHandlerFunc {
101+
// nats server info
102+
// Args:
103+
//
104+
// [<server>] Server ID or Name to inspect
105+
func (s *ServerTools) serverInfoHandler() server.ToolHandlerFunc {
92106
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
93-
output, err := n.executeNATSCommand("server", "check")
107+
var args []string
108+
args = append(args, "server", "info")
109+
110+
if server, ok := request.Params.Arguments["server"].(string); ok {
111+
args = append(args, server)
112+
}
113+
output, err := s.nats.GetSysExecutor().ExecuteCommand(args...)
94114
if err != nil {
95115
return nil, err
96116
}
97117
return mcp.NewToolResultText(output), nil
98118
}
99119
}
100120

101-
func (n *NATSServerTools) serverPingHandler() server.ToolHandlerFunc {
121+
// nats server ping
122+
// Args:
123+
//
124+
// [<expect>] How many servers to expect
125+
func (s *ServerTools) serverPingHandler() server.ToolHandlerFunc {
102126
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
103-
output, err := n.executeNATSCommand("server", "ping")
127+
var args []string
128+
args = append(args, "server", "ping")
129+
130+
if expect, ok := request.Params.Arguments["expect"].(string); ok {
131+
args = append(args, expect)
132+
}
133+
output, err := s.nats.GetSysExecutor().ExecuteCommand(args...)
104134
if err != nil {
105135
return nil, err
106136
}

0 commit comments

Comments
 (0)