Skip to content

Commit ac5c2ef

Browse files
ezynda3ezynda3
authored
Merge branch 'main' into update-arguments-parsing
2 parents a8474c7 + 077f546 commit ac5c2ef

23 files changed

+1106
-158
lines changed

‎client/client.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ func (c *Client) sendRequest(
104104

105105
request := transport.JSONRPCRequest{
106106
JSONRPC: mcp.JSONRPC_VERSION,
107-
ID: id,
107+
ID: mcp.NewRequestId(id),
108108
Method: method,
109109
Params: params,
110110
}

‎client/transport/interface.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,15 +27,15 @@ type Interface interface {
2727
}
2828

2929
type JSONRPCRequest struct {
30-
JSONRPC string `json:"jsonrpc"`
31-
ID int64 `json:"id"`
32-
Method string `json:"method"`
33-
Params any `json:"params,omitempty"`
30+
JSONRPC string `json:"jsonrpc"`
31+
ID mcp.RequestId `json:"id"`
32+
Method string `json:"method"`
33+
Params any `json:"params,omitempty"`
3434
}
3535

3636
type JSONRPCResponse struct {
3737
JSONRPC string `json:"jsonrpc"`
38-
ID *int64 `json:"id"`
38+
ID mcp.RequestId `json:"id"`
3939
Result json.RawMessage `json:"result"`
4040
Error *struct {
4141
Code int `json:"code"`

‎client/transport/sse.go

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ type SSE struct {
2525
baseURL *url.URL
2626
endpoint *url.URL
2727
httpClient *http.Client
28-
responses map[int64]chan *JSONRPCResponse
28+
responses map[string]chan *JSONRPCResponse
2929
mu sync.RWMutex
3030
onNotification func(mcp.JSONRPCNotification)
3131
notifyMu sync.RWMutex
@@ -62,7 +62,7 @@ func NewSSE(baseURL string, options ...ClientOption) (*SSE, error) {
6262
smc := &SSE{
6363
baseURL: parsedURL,
6464
httpClient: &http.Client{},
65-
responses: make(map[int64]chan *JSONRPCResponse),
65+
responses: make(map[string]chan *JSONRPCResponse),
6666
endpointChan: make(chan struct{}),
6767
headers: make(map[string]string),
6868
}
@@ -200,7 +200,7 @@ func (c *SSE) handleSSEEvent(event, data string) {
200200
}
201201

202202
// Handle notification
203-
if baseMessage.ID == nil {
203+
if baseMessage.ID.IsNil() {
204204
var notification mcp.JSONRPCNotification
205205
if err := json.Unmarshal([]byte(data), &notification); err != nil {
206206
return
@@ -213,14 +213,17 @@ func (c *SSE) handleSSEEvent(event, data string) {
213213
return
214214
}
215215

216+
// Create string key for map lookup
217+
idKey := baseMessage.ID.String()
218+
216219
c.mu.RLock()
217-
ch, ok := c.responses[*baseMessage.ID]
220+
ch, exists := c.responses[idKey]
218221
c.mu.RUnlock()
219222

220-
if ok {
223+
if exists {
221224
ch <- &baseMessage
222225
c.mu.Lock()
223-
delete(c.responses, *baseMessage.ID)
226+
delete(c.responses, idKey)
224227
c.mu.Unlock()
225228
}
226229
}
@@ -267,14 +270,17 @@ func (c *SSE) SendRequest(
267270
req.Header.Set(k, v)
268271
}
269272

273+
// Create string key for map lookup
274+
idKey := request.ID.String()
275+
270276
// Register response channel
271277
responseChan := make(chan *JSONRPCResponse, 1)
272278
c.mu.Lock()
273-
c.responses[request.ID] = responseChan
279+
c.responses[idKey] = responseChan
274280
c.mu.Unlock()
275281
deleteResponseChan := func() {
276282
c.mu.Lock()
277-
delete(c.responses, request.ID)
283+
delete(c.responses, idKey)
278284
c.mu.Unlock()
279285
}
280286

@@ -327,7 +333,7 @@ func (c *SSE) Close() error {
327333
for _, ch := range c.responses {
328334
close(ch)
329335
}
330-
c.responses = make(map[int64]chan *JSONRPCResponse)
336+
c.responses = make(map[string]chan *JSONRPCResponse)
331337
c.mu.Unlock()
332338

333339
return nil

‎client/transport/sse_test.go

Lines changed: 35 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ func TestSSE(t *testing.T) {
160160

161161
request := JSONRPCRequest{
162162
JSONRPC: "2.0",
163-
ID: 1,
163+
ID: mcp.NewRequestId(int64(1)),
164164
Method: "debug/echo",
165165
Params: params,
166166
}
@@ -174,7 +174,7 @@ func TestSSE(t *testing.T) {
174174
// Parse the result to verify echo
175175
var result struct {
176176
JSONRPC string `json:"jsonrpc"`
177-
ID int64 `json:"id"`
177+
ID mcp.RequestId `json:"id"`
178178
Method string `json:"method"`
179179
Params map[string]any `json:"params"`
180180
}
@@ -187,8 +187,11 @@ func TestSSE(t *testing.T) {
187187
if result.JSONRPC != "2.0" {
188188
t.Errorf("Expected JSONRPC value '2.0', got '%s'", result.JSONRPC)
189189
}
190-
if result.ID != 1 {
191-
t.Errorf("Expected ID 1, got %d", result.ID)
190+
idValue, ok := result.ID.Value().(int64)
191+
if !ok {
192+
t.Errorf("Expected ID to be int64, got %T", result.ID.Value())
193+
} else if idValue != 1 {
194+
t.Errorf("Expected ID 1, got %d", idValue)
192195
}
193196
if result.Method != "debug/echo" {
194197
t.Errorf("Expected method 'debug/echo', got '%s'", result.Method)
@@ -211,7 +214,7 @@ func TestSSE(t *testing.T) {
211214
// Prepare a request
212215
request := JSONRPCRequest{
213216
JSONRPC: "2.0",
214-
ID: 3,
217+
ID: mcp.NewRequestId(int64(3)),
215218
Method: "debug/echo",
216219
}
217220

@@ -292,7 +295,7 @@ func TestSSE(t *testing.T) {
292295
// Each request has a unique ID and payload
293296
request := JSONRPCRequest{
294297
JSONRPC: "2.0",
295-
ID: int64(100 + idx),
298+
ID: mcp.NewRequestId(int64(100 + idx)),
296299
Method: "debug/echo",
297300
Params: map[string]any{
298301
"requestIndex": idx,
@@ -317,15 +320,25 @@ func TestSSE(t *testing.T) {
317320
continue
318321
}
319322

320-
if responses[i] == nil || responses[i].ID == nil || *responses[i].ID != int64(100+i) {
321-
t.Errorf("Request %d: Expected ID %d, got %v", i, 100+i, responses[i])
323+
if responses[i] == nil {
324+
t.Errorf("Request %d: Response is nil", i)
325+
continue
326+
}
327+
328+
expectedId := int64(100 + i)
329+
idValue, ok := responses[i].ID.Value().(int64)
330+
if !ok {
331+
t.Errorf("Request %d: Expected ID to be int64, got %T", i, responses[i].ID.Value())
332+
continue
333+
} else if idValue != expectedId {
334+
t.Errorf("Request %d: Expected ID %d, got %d", i, expectedId, idValue)
322335
continue
323336
}
324337

325338
// Parse the result to verify echo
326339
var result struct {
327340
JSONRPC string `json:"jsonrpc"`
328-
ID int64 `json:"id"`
341+
ID mcp.RequestId `json:"id"`
329342
Method string `json:"method"`
330343
Params map[string]any `json:"params"`
331344
}
@@ -336,8 +349,11 @@ func TestSSE(t *testing.T) {
336349
}
337350

338351
// Verify data matches what was sent
339-
if result.ID != int64(100+i) {
340-
t.Errorf("Request %d: Expected echoed ID %d, got %d", i, 100+i, result.ID)
352+
idValue, ok = result.ID.Value().(int64)
353+
if !ok {
354+
t.Errorf("Request %d: Expected ID to be int64, got %T", i, result.ID.Value())
355+
} else if idValue != int64(100+i) {
356+
t.Errorf("Request %d: Expected echoed ID %d, got %d", i, 100+i, idValue)
341357
}
342358

343359
if result.Method != "debug/echo" {
@@ -356,7 +372,7 @@ func TestSSE(t *testing.T) {
356372
// Prepare a request
357373
request := JSONRPCRequest{
358374
JSONRPC: "2.0",
359-
ID: 100,
375+
ID: mcp.NewRequestId(int64(100)),
360376
Method: "debug/echo_error_string",
361377
}
362378

@@ -378,8 +394,11 @@ func TestSSE(t *testing.T) {
378394
if responseError.Method != "debug/echo_error_string" {
379395
t.Errorf("Expected method 'debug/echo_error_string', got '%s'", responseError.Method)
380396
}
381-
if responseError.ID != 100 {
382-
t.Errorf("Expected ID 100, got %d", responseError.ID)
397+
idValue, ok := responseError.ID.Value().(int64)
398+
if !ok {
399+
t.Errorf("Expected ID to be int64, got %T", responseError.ID.Value())
400+
} else if idValue != 100 {
401+
t.Errorf("Expected ID 100, got %d", idValue)
383402
}
384403
if responseError.JSONRPC != "2.0" {
385404
t.Errorf("Expected JSONRPC '2.0', got '%s'", responseError.JSONRPC)
@@ -453,7 +472,7 @@ func TestSSEErrors(t *testing.T) {
453472
// Prepare a request
454473
request := JSONRPCRequest{
455474
JSONRPC: "2.0",
456-
ID: 99,
475+
ID: mcp.NewRequestId(int64(99)),
457476
Method: "ping",
458477
}
459478

@@ -492,7 +511,7 @@ func TestSSEErrors(t *testing.T) {
492511
// Try to send a request after close
493512
request := JSONRPCRequest{
494513
JSONRPC: "2.0",
495-
ID: 1,
514+
ID: mcp.NewRequestId(int64(1)),
496515
Method: "ping",
497516
}
498517

‎client/transport/stdio.go

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ type Stdio struct {
2626
stdin io.WriteCloser
2727
stdout *bufio.Reader
2828
stderr io.ReadCloser
29-
responses map[int64]chan *JSONRPCResponse
29+
responses map[string]chan *JSONRPCResponse
3030
mu sync.RWMutex
3131
done chan struct{}
3232
onNotification func(mcp.JSONRPCNotification)
@@ -42,7 +42,7 @@ func NewIO(input io.Reader, output io.WriteCloser, logging io.ReadCloser) *Stdio
4242
stdout: bufio.NewReader(input),
4343
stderr: logging,
4444

45-
responses: make(map[int64]chan *JSONRPCResponse),
45+
responses: make(map[string]chan *JSONRPCResponse),
4646
done: make(chan struct{}),
4747
}
4848
}
@@ -61,7 +61,7 @@ func NewStdio(
6161
args: args,
6262
env: env,
6363

64-
responses: make(map[int64]chan *JSONRPCResponse),
64+
responses: make(map[string]chan *JSONRPCResponse),
6565
done: make(chan struct{}),
6666
}
6767

@@ -181,7 +181,7 @@ func (c *Stdio) readResponses() {
181181
}
182182

183183
// Handle notification
184-
if baseMessage.ID == nil {
184+
if baseMessage.ID.IsNil() {
185185
var notification mcp.JSONRPCNotification
186186
if err := json.Unmarshal([]byte(line), &notification); err != nil {
187187
continue
@@ -194,14 +194,17 @@ func (c *Stdio) readResponses() {
194194
continue
195195
}
196196

197+
// Create string key for map lookup
198+
idKey := baseMessage.ID.String()
199+
197200
c.mu.RLock()
198-
ch, ok := c.responses[*baseMessage.ID]
201+
ch, exists := c.responses[idKey]
199202
c.mu.RUnlock()
200203

201-
if ok {
204+
if exists {
202205
ch <- &baseMessage
203206
c.mu.Lock()
204-
delete(c.responses, *baseMessage.ID)
207+
delete(c.responses, idKey)
205208
c.mu.Unlock()
206209
}
207210
}
@@ -227,14 +230,17 @@ func (c *Stdio) SendRequest(
227230
}
228231
requestBytes = append(requestBytes, '\n')
229232

233+
// Create string key for map lookup
234+
idKey := request.ID.String()
235+
230236
// Register response channel
231237
responseChan := make(chan *JSONRPCResponse, 1)
232238
c.mu.Lock()
233-
c.responses[request.ID] = responseChan
239+
c.responses[idKey] = responseChan
234240
c.mu.Unlock()
235241
deleteResponseChan := func() {
236242
c.mu.Lock()
237-
delete(c.responses, request.ID)
243+
delete(c.responses, idKey)
238244
c.mu.Unlock()
239245
}
240246

0 commit comments

Comments
 (0)