Fix LSP exit status after shutdown

tsc/internal/lsp/server.go

@@ -113,6 +113,8 @@ type Writer interface {
 	Write(msg *lsproto.Message) error
 }
 
+var errExitBeforeShutdown = errors.New("exit notification received before shutdown")
+
 type lspReader struct {
 	r *lsproto.BaseReader
 }
@@ -178,6 +180,7 @@ type Server struct {
 
 	logger                  *logger
 	initStarted             atomic.Bool
+	shutdownRequested       atomic.Bool
 	clientSeq               atomic.Int32
 	requestQueue            *dynamicQueue[*lsproto.RequestMessage]
 	outgoingQueue           *dynamicQueue[*lsproto.Message]
@@ -936,6 +939,13 @@ func (s *Server) readLoop(ctx context.Context) error {
 				}
 				continue
 			}
+			if errors.Is(err, io.EOF) {
+				// Queue EOF after pending requests so transport closure cannot preempt dispatch.
+				if err := s.requestQueue.Put(ctx, nil); err != nil {
+					return err
+				}
+				return nil
+			}
 			return err
 		}
 
@@ -1008,8 +1018,14 @@ func (s *Server) dispatchLoop(ctx context.Context) error {
 	for {
 		req, err := s.requestQueue.Get(ctx)
 		if err != nil {
+			if cause := context.Cause(ctx); errors.Is(cause, io.EOF) || errors.Is(cause, errExitBeforeShutdown) {
+				return cause
+			}
 			return err
 		}
+		if req == nil {
+			return io.EOF
+		}
 
 		s.lastRequestTimeMs.Store(time.Now().UnixMilli())
 		requestCtx := locale.WithLocale(ctx, s.GetLocale())
@@ -1029,8 +1045,8 @@ func (s *Server) dispatchLoop(ctx context.Context) error {
 				if err := s.sendError(req.ID, lsproto.ErrorCodeRequestCancelled); err != nil {
 					lspExit(err)
 				}
-			} else if errors.Is(err, io.EOF) {
-				lspExit(nil)
+			} else if errors.Is(err, io.EOF) || errors.Is(err, errExitBeforeShutdown) {
+				lspExit(err)
 			} else {
 				if err := s.sendError(req.ID, err); err != nil {
 					lspExit(err)
@@ -1336,7 +1352,7 @@ var handlers = sync.OnceValue(func() handlerMap {
 
 func (handlers handlerMap) registerNotificationHandler[Req any](info lsproto.NotificationInfo[Req], fn func(*Server, context.Context, Req) error) {
 	handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) (func() error, error) {
-		if s.session == nil && req.Method != lsproto.MethodInitialized {
+		if s.session == nil && req.Method != lsproto.MethodInitialized && req.Method != lsproto.MethodExit {
 			return nil, lsproto.ErrorCodeServerNotInitialized
 		}
 
@@ -1834,10 +1850,14 @@ func (s *Server) handleShutdown(ctx context.Context, _ lsproto.NoParams, _ *lspr
 	}
 	s.closeAPISessions()
 	s.session.Close()
+	s.shutdownRequested.Store(true)
 	return lsproto.ShutdownResponse{}, nil
 }
 
 func (s *Server) handleExit(ctx context.Context, _ lsproto.NoParams) error {
+	if !s.shutdownRequested.Load() {
+		return errExitBeforeShutdown
+	}
 	return io.EOF
 }
 

tsc/internal/lsp/server_test.go

@@ -23,6 +23,16 @@ type shutdownTestWriter struct{}
 
 func (shutdownTestWriter) Write(*lsproto.Message) error { return nil }
 
+type exitThenEOFTestReader struct{ read bool }
+
+func (r *exitThenEOFTestReader) Read() (*lsproto.Message, error) {
+	if !r.read {
+		r.read = true
+		return (&lsproto.RequestMessage{Method: lsproto.MethodExit}).Message(), nil
+	}
+	return nil, io.EOF
+}
+
 type cancelingTestWriter struct {
 	writer   Writer
 	cancel   context.CancelFunc
@@ -72,6 +82,78 @@ func TestGetContentMapperWorkspaceRoots(t *testing.T) {
 	})
 }
 
+func TestDispatchLoopReturnsEOFOnExit(t *testing.T) {
+	t.Parallel()
+
+	server := NewServer(&ServerOptions{
+		In:  shutdownTestReader{},
+		Out: shutdownTestWriter{},
+		Err: io.Discard,
+		Cwd: "/test",
+	})
+	server.session = &project.Session{}
+	server.shutdownRequested.Store(true)
+	if err := server.requestQueue.Put(t.Context(), &lsproto.RequestMessage{Method: lsproto.MethodExit}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := server.dispatchLoop(t.Context()); !errors.Is(err, io.EOF) {
+		t.Fatalf("dispatchLoop() error = %v, want %v", err, io.EOF)
+	}
+}
+
+func TestDispatchLoopReturnsErrorWhenExitBeforeShutdown(t *testing.T) {
+	t.Parallel()
+
+	server := NewServer(&ServerOptions{
+		In:  shutdownTestReader{},
+		Out: shutdownTestWriter{},
+		Err: io.Discard,
+		Cwd: "/test",
+	})
+	if err := server.requestQueue.Put(t.Context(), &lsproto.RequestMessage{Method: lsproto.MethodExit}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := server.dispatchLoop(t.Context()); !errors.Is(err, errExitBeforeShutdown) {
+		t.Fatalf("dispatchLoop() error = %v, want %v", err, errExitBeforeShutdown)
+	}
+}
+
+func TestRunExitsSuccessfullyAfterShutdown(t *testing.T) {
+	t.Parallel()
+
+	reader := &exitThenEOFTestReader{}
+	server := NewServer(&ServerOptions{
+		In:  reader,
+		Out: shutdownTestWriter{},
+		Err: io.Discard,
+		Cwd: "/test",
+	})
+	server.session = &project.Session{}
+	server.shutdownRequested.Store(true)
+
+	err := server.Run(t.Context())
+	if err != nil {
+		t.Fatalf("Run() error = %v, want nil", err)
+	}
+}
+
+func TestRunReturnsErrorWhenExitPrecedesShutdown(t *testing.T) {
+	t.Parallel()
+
+	server := NewServer(&ServerOptions{
+		In:  &exitThenEOFTestReader{},
+		Out: shutdownTestWriter{},
+		Err: io.Discard,
+		Cwd: "/test",
+	})
+
+	if err := server.Run(t.Context()); !errors.Is(err, errExitBeforeShutdown) {
+		t.Fatalf("Run() error = %v, want %v", err, errExitBeforeShutdown)
+	}
+}
+
 // TestServerShutdownNoDeadlock verifies that operations after shutdown
 // don't block.
 func TestServerShutdownNoDeadlock(t *testing.T) {