1 files changed, 41 insertions, 0 deletions
diff --git a/internal/lsp/init_shutdown_test.go b/internal/lsp/init_shutdown_test.go
index 2847170..4e1bd2f 100644
--- a/internal/lsp/init_shutdown_test.go
+++ b/internal/lsp/init_shutdown_test.go
@@ -2,10 +2,13 @@ package lsp
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"io"
 	"log"
 	"testing"
+	"time"
 )
 
 func TestHandleShutdown_Replies(t *testing.T) {
@@ -20,3 +23,41 @@ func TestHandleShutdown_Replies(t *testing.T) {
 		t.Fatalf("unexpected shutdown response: %+v", resp)
 	}
 }
+
+func TestHandleShutdown_CancelsServerContext(t *testing.T) {
+	var out bytes.Buffer
+	s := NewServer(bytes.NewReader(nil), &out, log.New(io.Discard, "", 0), ServerOptions{})
+	req := Request{JSONRPC: "2.0", ID: json.RawMessage("12"), Method: "shutdown"}
+	s.handleShutdown(req)
+
+	ctx, cancel := s.requestTimeoutContext(2 * time.Second)
+	defer cancel()
+	select {
+	case <-ctx.Done():
+		if !errors.Is(ctx.Err(), context.Canceled) {
+			t.Fatalf("expected canceled context, got %v", ctx.Err())
+		}
+	default:
+		t.Fatalf("expected canceled context after shutdown")
+	}
+}
+
+func TestHandleExit_CancelsServerContext(t *testing.T) {
+	var out bytes.Buffer
+	s := NewServer(bytes.NewReader(nil), &out, log.New(io.Discard, "", 0), ServerOptions{})
+	s.handleExit()
+	if !s.exited.Load() {
+		t.Fatalf("expected exited flag to be set")
+	}
+
+	ctx, cancel := s.requestTimeoutContext(2 * time.Second)
+	defer cancel()
+	select {
+	case <-ctx.Done():
+		if !errors.Is(ctx.Err(), context.Canceled) {
+			t.Fatalf("expected canceled context, got %v", ctx.Err())
+		}
+	default:
+		t.Fatalf("expected canceled context after exit")
+	}
+}