4
4
"context"
5
5
"encoding/json"
6
6
"errors"
7
+ "io"
8
+ "strings"
7
9
"sync"
8
10
"testing"
9
11
"time"
@@ -15,6 +17,39 @@ import (
15
17
"github.com/mark3labs/mcp-go/mcp"
16
18
)
17
19
20
+ // mockReaderWithError is a mock io.ReadCloser that simulates reading some data
21
+ // and then returning a specific error
22
+ type mockReaderWithError struct {
23
+ data []byte
24
+ err error
25
+ position int
26
+ closed bool
27
+ }
28
+
29
+ func (m * mockReaderWithError ) Read (p []byte ) (n int , err error ) {
30
+ if m .closed {
31
+ return 0 , io .EOF
32
+ }
33
+
34
+ if m .position >= len (m .data ) {
35
+ return 0 , m .err
36
+ }
37
+
38
+ n = copy (p , m .data [m .position :])
39
+ m .position += n
40
+
41
+ if m .position >= len (m .data ) {
42
+ return n , m .err
43
+ }
44
+
45
+ return n , nil
46
+ }
47
+
48
+ func (m * mockReaderWithError ) Close () error {
49
+ m .closed = true
50
+ return nil
51
+ }
52
+
18
53
// startMockSSEEchoServer starts a test HTTP server that implements
19
54
// a minimal SSE-based echo server for testing purposes.
20
55
// It returns the server URL and a function to close the server.
@@ -508,6 +543,218 @@ func TestSSE(t *testing.T) {
508
543
}
509
544
})
510
545
546
+ t .Run ("NO_ERROR_WithoutConnectionLostHandler" , func (t * testing.T ) {
547
+ // Test that NO_ERROR without connection lost handler maintains backward compatibility
548
+ // When no connection lost handler is set, NO_ERROR should be treated as a regular error
549
+
550
+ // Create a mock Reader that simulates NO_ERROR
551
+ mockReader := & mockReaderWithError {
552
+ data : []byte ("event: endpoint\n data: /message\n \n " ),
553
+ err : errors .New ("connection closed: NO_ERROR" ),
554
+ }
555
+
556
+ // Create SSE transport
557
+ url , closeF := startMockSSEEchoServer ()
558
+ defer closeF ()
559
+
560
+ trans , err := NewSSE (url )
561
+ if err != nil {
562
+ t .Fatal (err )
563
+ }
564
+
565
+ // DO NOT set connection lost handler to test backward compatibility
566
+
567
+ // Capture stderr to verify the error is printed (backward compatible behavior)
568
+ // Since we can't easily capture fmt.Printf output in tests, we'll just verify
569
+ // that the readSSE method returns without calling any handler
570
+
571
+ // Directly test the readSSE method with our mock reader
572
+ go trans .readSSE (mockReader )
573
+
574
+ // Wait for readSSE to complete
575
+ time .Sleep (100 * time .Millisecond )
576
+
577
+ // The test passes if readSSE completes without panicking or hanging
578
+ // In backward compatibility mode, NO_ERROR should be treated as a regular error
579
+ t .Log ("Backward compatibility test passed: NO_ERROR handled as regular error when no handler is set" )
580
+ })
581
+
582
+ t .Run ("NO_ERROR_ConnectionLost" , func (t * testing.T ) {
583
+ // Test that NO_ERROR in HTTP/2 connection loss is properly handled
584
+ // This test verifies that when a connection is lost in a way that produces
585
+ // an error message containing "NO_ERROR", the connection lost handler is called
586
+
587
+ var connectionLostCalled bool
588
+ var connectionLostError error
589
+ var mu sync.Mutex
590
+
591
+ // Create a mock Reader that simulates connection loss with NO_ERROR
592
+ mockReader := & mockReaderWithError {
593
+ data : []byte ("event: endpoint\n data: /message\n \n " ),
594
+ err : errors .New ("http2: stream closed with error code NO_ERROR" ),
595
+ }
596
+
597
+ // Create SSE transport
598
+ url , closeF := startMockSSEEchoServer ()
599
+ defer closeF ()
600
+
601
+ trans , err := NewSSE (url )
602
+ if err != nil {
603
+ t .Fatal (err )
604
+ }
605
+
606
+ // Set connection lost handler
607
+ trans .SetConnectionLostHandler (func (err error ) {
608
+ mu .Lock ()
609
+ defer mu .Unlock ()
610
+ connectionLostCalled = true
611
+ connectionLostError = err
612
+ })
613
+
614
+ // Directly test the readSSE method with our mock reader that simulates NO_ERROR
615
+ go trans .readSSE (mockReader )
616
+
617
+ // Wait for connection lost handler to be called
618
+ timeout := time .After (1 * time .Second )
619
+ ticker := time .NewTicker (10 * time .Millisecond )
620
+ defer ticker .Stop ()
621
+
622
+ for {
623
+ select {
624
+ case <- timeout :
625
+ t .Fatal ("Connection lost handler was not called within timeout for NO_ERROR connection loss" )
626
+ case <- ticker .C :
627
+ mu .Lock ()
628
+ called := connectionLostCalled
629
+ err := connectionLostError
630
+ mu .Unlock ()
631
+
632
+ if called {
633
+ if err == nil {
634
+ t .Fatal ("Expected connection lost error, got nil" )
635
+ }
636
+
637
+ // Verify that the error contains "NO_ERROR" string
638
+ if ! strings .Contains (err .Error (), "NO_ERROR" ) {
639
+ t .Errorf ("Expected error to contain 'NO_ERROR', got: %v" , err )
640
+ }
641
+
642
+ t .Logf ("Connection lost handler called with NO_ERROR: %v" , err )
643
+ return
644
+ }
645
+ }
646
+ }
647
+ })
648
+
649
+ t .Run ("NO_ERROR_Handling" , func (t * testing.T ) {
650
+ // Test specific NO_ERROR string handling in readSSE method
651
+ // This tests the code path at line 209 where NO_ERROR is checked
652
+
653
+ // Create a mock Reader that simulates an error containing "NO_ERROR"
654
+ mockReader := & mockReaderWithError {
655
+ data : []byte ("event: endpoint\n data: /message\n \n " ),
656
+ err : errors .New ("connection closed: NO_ERROR" ),
657
+ }
658
+
659
+ // Create SSE transport
660
+ url , closeF := startMockSSEEchoServer ()
661
+ defer closeF ()
662
+
663
+ trans , err := NewSSE (url )
664
+ if err != nil {
665
+ t .Fatal (err )
666
+ }
667
+
668
+ var connectionLostCalled bool
669
+ var connectionLostError error
670
+ var mu sync.Mutex
671
+
672
+ // Set connection lost handler to verify it's called for NO_ERROR
673
+ trans .SetConnectionLostHandler (func (err error ) {
674
+ mu .Lock ()
675
+ defer mu .Unlock ()
676
+ connectionLostCalled = true
677
+ connectionLostError = err
678
+ })
679
+
680
+ // Directly test the readSSE method with our mock reader
681
+ go trans .readSSE (mockReader )
682
+
683
+ // Wait for connection lost handler to be called
684
+ timeout := time .After (1 * time .Second )
685
+ ticker := time .NewTicker (10 * time .Millisecond )
686
+ defer ticker .Stop ()
687
+
688
+ for {
689
+ select {
690
+ case <- timeout :
691
+ t .Fatal ("Connection lost handler was not called within timeout for NO_ERROR" )
692
+ case <- ticker .C :
693
+ mu .Lock ()
694
+ called := connectionLostCalled
695
+ err := connectionLostError
696
+ mu .Unlock ()
697
+
698
+ if called {
699
+ if err == nil {
700
+ t .Fatal ("Expected connection lost error with NO_ERROR, got nil" )
701
+ }
702
+
703
+ // Verify that the error contains "NO_ERROR" string
704
+ if ! strings .Contains (err .Error (), "NO_ERROR" ) {
705
+ t .Errorf ("Expected error to contain 'NO_ERROR', got: %v" , err )
706
+ }
707
+
708
+ t .Logf ("Successfully handled NO_ERROR: %v" , err )
709
+ return
710
+ }
711
+ }
712
+ }
713
+ })
714
+
715
+ t .Run ("RegularError_DoesNotTriggerConnectionLost" , func (t * testing.T ) {
716
+ // Test that regular errors (not containing NO_ERROR) do not trigger connection lost handler
717
+
718
+ // Create a mock Reader that simulates a regular error
719
+ mockReader := & mockReaderWithError {
720
+ data : []byte ("event: endpoint\n data: /message\n \n " ),
721
+ err : errors .New ("regular connection error" ),
722
+ }
723
+
724
+ // Create SSE transport
725
+ url , closeF := startMockSSEEchoServer ()
726
+ defer closeF ()
727
+
728
+ trans , err := NewSSE (url )
729
+ if err != nil {
730
+ t .Fatal (err )
731
+ }
732
+
733
+ var connectionLostCalled bool
734
+ var mu sync.Mutex
735
+
736
+ // Set connection lost handler - this should NOT be called for regular errors
737
+ trans .SetConnectionLostHandler (func (err error ) {
738
+ mu .Lock ()
739
+ defer mu .Unlock ()
740
+ connectionLostCalled = true
741
+ })
742
+
743
+ // Directly test the readSSE method with our mock reader
744
+ go trans .readSSE (mockReader )
745
+
746
+ // Wait and verify connection lost handler is NOT called
747
+ time .Sleep (200 * time .Millisecond )
748
+
749
+ mu .Lock ()
750
+ called := connectionLostCalled
751
+ mu .Unlock ()
752
+
753
+ if called {
754
+ t .Error ("Connection lost handler should not be called for regular errors" )
755
+ }
756
+ })
757
+
511
758
}
512
759
513
760
func TestSSEErrors (t * testing.T ) {
0 commit comments