diff --git a/internal/lsp/cmd/check.go b/internal/lsp/cmd/check.go
index 4fdcff6..3231828 100644
--- a/internal/lsp/cmd/check.go
+++ b/internal/lsp/cmd/check.go
@@ -8,8 +8,8 @@
 	"context"
 	"flag"
 	"fmt"
+	"time"
 
-	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/span"
 )
 
@@ -18,16 +18,6 @@
 	app *Application
 }
 
-type checkClient struct {
-	baseClient
-	diagnostics chan entry
-}
-
-type entry struct {
-	uri         span.URI
-	diagnostics []protocol.Diagnostic
-}
-
 func (c *check) Name() string      { return "check" }
 func (c *check) Usage() string     { return "<filename>" }
 func (c *check) ShortHelp() string { return "show diagnostic results for the specified file" }
@@ -49,49 +39,35 @@
 		// no files, so no results
 		return nil
 	}
-	client := &checkClient{
-		diagnostics: make(chan entry),
-	}
-	checking := map[span.URI]*protocol.ColumnMapper{}
+	checking := map[span.URI]*cmdFile{}
 	// now we ready to kick things off
-	_, err := c.app.connect(ctx, client)
+	conn, err := c.app.connect(ctx)
 	if err != nil {
 		return err
 	}
 	for _, arg := range args {
 		uri := span.FileURI(arg)
-		m, err := client.AddFile(ctx, uri)
-		if err != nil {
-			return err
+		file := conn.AddFile(ctx, uri)
+		if file.err != nil {
+			return file.err
 		}
-		checking[uri] = m
+		checking[uri] = file
 	}
 	// now wait for results
-	for entry := range client.diagnostics {
-		//TODO:timeout?
-		m, found := checking[entry.uri]
-		if !found {
-			continue
+	//TODO: maybe conn.ExecuteCommand(ctx, &protocol.ExecuteCommandParams{Command: "gopls-wait-idle"})
+	for _, file := range checking {
+		select {
+		case <-file.hasDiagnostics:
+		case <-time.Tick(30 * time.Second):
+			return fmt.Errorf("timed out waiting for results from %v", file.uri)
 		}
-		for _, d := range entry.diagnostics {
-			spn, err := m.RangeSpan(d.Range)
+		for _, d := range file.diagnostics {
+			spn, err := file.mapper.RangeSpan(d.Range)
 			if err != nil {
 				return fmt.Errorf("Could not convert position %v for %q", d.Range, d.Message)
 			}
 			fmt.Printf("%v: %v\n", spn, d.Message)
 		}
-		delete(checking, entry.uri)
-		if len(checking) == 0 {
-			return nil
-		}
-	}
-	return fmt.Errorf("did not get all results")
-}
-
-func (c *checkClient) PublishDiagnostics(ctx context.Context, p *protocol.PublishDiagnosticsParams) error {
-	c.diagnostics <- entry{
-		uri:         span.URI(p.URI),
-		diagnostics: p.Diagnostics,
 	}
 	return nil
 }
diff --git a/internal/lsp/cmd/check_test.go b/internal/lsp/cmd/check_test.go
index 4efd7d1..fad0d78 100644
--- a/internal/lsp/cmd/check_test.go
+++ b/internal/lsp/cmd/check_test.go
@@ -7,7 +7,6 @@
 import (
 	"context"
 	"fmt"
-	"runtime"
 	"strings"
 	"testing"
 
@@ -17,9 +16,6 @@
 )
 
 func (r *runner) Diagnostics(t *testing.T, data tests.Diagnostics) {
-	if runtime.GOOS != "linux" || isRace {
-		t.Skip("currently uses too much memory, see issue #31611")
-	}
 	for uri, want := range data {
 		if len(want) == 1 && want[0].Message == "" {
 			continue
@@ -28,8 +24,7 @@
 		if err != nil {
 			t.Fatal(err)
 		}
-		args := []string{"-remote=internal"}
-		args = append(args, "check", fname)
+		args := []string{"-remote=internal", "check", fname}
 		out := captureStdOut(t, func() {
 			tool.Main(context.Background(), r.app, args)
 		})
diff --git a/internal/lsp/cmd/cmd.go b/internal/lsp/cmd/cmd.go
index af05dcd..e4a8da5 100644
--- a/internal/lsp/cmd/cmd.go
+++ b/internal/lsp/cmd/cmd.go
@@ -19,6 +19,7 @@
 	"net"
 	"os"
 	"strings"
+	"sync"
 
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/internal/jsonrpc2"
@@ -120,62 +121,104 @@
 	}
 }
 
-type cmdClient interface {
-	protocol.Client
+var (
+	internalMu          sync.Mutex
+	internalConnections = make(map[string]*connection)
+)
 
-	prepare(app *Application, server protocol.Server)
-}
-
-func (app *Application) connect(ctx context.Context, client cmdClient) (protocol.Server, error) {
-	var server protocol.Server
+func (app *Application) connect(ctx context.Context) (*connection, error) {
 	switch app.Remote {
 	case "":
-		server = lsp.NewClientServer(client)
+		connection := newConnection(app)
+		connection.Server = lsp.NewClientServer(connection.Client)
+		return connection, connection.initialize(ctx)
 	case "internal":
+		internalMu.Lock()
+		defer internalMu.Unlock()
+		if c := internalConnections[app.Config.Dir]; c != nil {
+			return c, nil
+		}
+		connection := newConnection(app)
+		ctx := context.Background() //TODO:a way of shutting down the internal server
 		cr, sw, _ := os.Pipe()
 		sr, cw, _ := os.Pipe()
 		var jc *jsonrpc2.Conn
-		jc, server, _ = protocol.NewClient(jsonrpc2.NewHeaderStream(cr, cw), client)
+		jc, connection.Server, _ = protocol.NewClient(jsonrpc2.NewHeaderStream(cr, cw), connection.Client)
 		go jc.Run(ctx)
 		go lsp.NewServer(jsonrpc2.NewHeaderStream(sr, sw)).Run(ctx)
+		if err := connection.initialize(ctx); err != nil {
+			return nil, err
+		}
+		internalConnections[app.Config.Dir] = connection
+		return connection, nil
 	default:
+		connection := newConnection(app)
 		conn, err := net.Dial("tcp", app.Remote)
 		if err != nil {
 			return nil, err
 		}
 		stream := jsonrpc2.NewHeaderStream(conn, conn)
 		var jc *jsonrpc2.Conn
-		jc, server, _ = protocol.NewClient(stream, client)
+		jc, connection.Server, _ = protocol.NewClient(stream, connection.Client)
 		go jc.Run(ctx)
+		return connection, connection.initialize(ctx)
 	}
+}
 
+func (c *connection) initialize(ctx context.Context) error {
 	params := &protocol.InitializeParams{}
-	params.RootURI = string(span.FileURI(app.Config.Dir))
+	params.RootURI = string(span.FileURI(c.Client.app.Config.Dir))
 	params.Capabilities.Workspace.Configuration = true
 	params.Capabilities.TextDocument.Hover.ContentFormat = []protocol.MarkupKind{protocol.PlainText}
-
-	client.prepare(app, server)
-	if _, err := server.Initialize(ctx, params); err != nil {
-		return nil, err
+	if _, err := c.Server.Initialize(ctx, params); err != nil {
+		return err
 	}
-	if err := server.Initialized(ctx, &protocol.InitializedParams{}); err != nil {
-		return nil, err
+	if err := c.Server.Initialized(ctx, &protocol.InitializedParams{}); err != nil {
+		return err
 	}
-	return server, nil
+	return nil
 }
 
-type baseClient struct {
+type connection struct {
 	protocol.Server
-	app    *Application
-	server protocol.Server
-	fset   *token.FileSet
+	Client *cmdClient
 }
 
-func (c *baseClient) ShowMessage(ctx context.Context, p *protocol.ShowMessageParams) error { return nil }
-func (c *baseClient) ShowMessageRequest(ctx context.Context, p *protocol.ShowMessageRequestParams) (*protocol.MessageActionItem, error) {
+type cmdClient struct {
+	protocol.Server
+	app  *Application
+	fset *token.FileSet
+
+	filesMu sync.Mutex
+	files   map[span.URI]*cmdFile
+}
+
+type cmdFile struct {
+	uri            span.URI
+	mapper         *protocol.ColumnMapper
+	err            error
+	added          bool
+	hasDiagnostics chan struct{}
+	diagnostics    []protocol.Diagnostic
+}
+
+func newConnection(app *Application) *connection {
+	return &connection{
+		Client: &cmdClient{
+			app:   app,
+			fset:  token.NewFileSet(),
+			files: make(map[span.URI]*cmdFile),
+		},
+	}
+}
+
+func (c *cmdClient) ShowMessage(ctx context.Context, p *protocol.ShowMessageParams) error { return nil }
+
+func (c *cmdClient) ShowMessageRequest(ctx context.Context, p *protocol.ShowMessageRequestParams) (*protocol.MessageActionItem, error) {
 	return nil, nil
 }
-func (c *baseClient) LogMessage(ctx context.Context, p *protocol.LogMessageParams) error {
+
+func (c *cmdClient) LogMessage(ctx context.Context, p *protocol.LogMessageParams) error {
 	switch p.Type {
 	case protocol.Error:
 		log.Print("Error:", p.Message)
@@ -190,17 +233,22 @@
 	}
 	return nil
 }
-func (c *baseClient) Event(ctx context.Context, t *interface{}) error { return nil }
-func (c *baseClient) RegisterCapability(ctx context.Context, p *protocol.RegistrationParams) error {
+
+func (c *cmdClient) Event(ctx context.Context, t *interface{}) error { return nil }
+
+func (c *cmdClient) RegisterCapability(ctx context.Context, p *protocol.RegistrationParams) error {
 	return nil
 }
-func (c *baseClient) UnregisterCapability(ctx context.Context, p *protocol.UnregistrationParams) error {
+
+func (c *cmdClient) UnregisterCapability(ctx context.Context, p *protocol.UnregistrationParams) error {
 	return nil
 }
-func (c *baseClient) WorkspaceFolders(ctx context.Context) ([]protocol.WorkspaceFolder, error) {
+
+func (c *cmdClient) WorkspaceFolders(ctx context.Context) ([]protocol.WorkspaceFolder, error) {
 	return nil, nil
 }
-func (c *baseClient) Configuration(ctx context.Context, p *protocol.ConfigurationParams) ([]interface{}, error) {
+
+func (c *cmdClient) Configuration(ctx context.Context, p *protocol.ConfigurationParams) ([]interface{}, error) {
 	results := make([]interface{}, len(p.Items))
 	for i, item := range p.Items {
 		if item.Section != "gopls" {
@@ -218,36 +266,66 @@
 	}
 	return results, nil
 }
-func (c *baseClient) ApplyEdit(ctx context.Context, p *protocol.ApplyWorkspaceEditParams) (*protocol.ApplyWorkspaceEditResponse, error) {
+
+func (c *cmdClient) ApplyEdit(ctx context.Context, p *protocol.ApplyWorkspaceEditParams) (*protocol.ApplyWorkspaceEditResponse, error) {
 	return &protocol.ApplyWorkspaceEditResponse{Applied: false, FailureReason: "not implemented"}, nil
 }
-func (c *baseClient) PublishDiagnostics(ctx context.Context, p *protocol.PublishDiagnosticsParams) error {
+
+func (c *cmdClient) PublishDiagnostics(ctx context.Context, p *protocol.PublishDiagnosticsParams) error {
+	c.filesMu.Lock()
+	defer c.filesMu.Unlock()
+	uri := span.URI(p.URI)
+	file := c.getFile(ctx, uri)
+	hadDiagnostics := file.diagnostics != nil
+	file.diagnostics = p.Diagnostics
+	if file.diagnostics == nil {
+		file.diagnostics = []protocol.Diagnostic{}
+	}
+	if !hadDiagnostics {
+		close(file.hasDiagnostics)
+	}
 	return nil
 }
 
-func (c *baseClient) prepare(app *Application, server protocol.Server) {
-	c.app = app
-	c.server = server
-	c.fset = token.NewFileSet()
+func (c *cmdClient) getFile(ctx context.Context, uri span.URI) *cmdFile {
+	file, found := c.files[uri]
+	if !found {
+		file = &cmdFile{
+			uri:            uri,
+			hasDiagnostics: make(chan struct{}),
+		}
+		c.files[uri] = file
+	}
+	if file.mapper == nil {
+		fname, err := uri.Filename()
+		if err != nil {
+			file.err = fmt.Errorf("%v: %v", uri, err)
+			return file
+		}
+		content, err := ioutil.ReadFile(fname)
+		if err != nil {
+			file.err = fmt.Errorf("%v: %v", uri, err)
+			return file
+		}
+		f := c.fset.AddFile(fname, -1, len(content))
+		f.SetLinesForContent(content)
+		file.mapper = protocol.NewColumnMapper(uri, c.fset, f, content)
+	}
+	return file
 }
 
-func (c *baseClient) AddFile(ctx context.Context, uri span.URI) (*protocol.ColumnMapper, error) {
-	fname, err := uri.Filename()
-	if err != nil {
-		return nil, fmt.Errorf("%v: %v", uri, err)
+func (c *connection) AddFile(ctx context.Context, uri span.URI) *cmdFile {
+	c.Client.filesMu.Lock()
+	defer c.Client.filesMu.Unlock()
+	file := c.Client.getFile(ctx, uri)
+	if !file.added {
+		file.added = true
+		p := &protocol.DidOpenTextDocumentParams{}
+		p.TextDocument.URI = string(uri)
+		p.TextDocument.Text = string(file.mapper.Content)
+		if err := c.Server.DidOpen(ctx, p); err != nil {
+			file.err = fmt.Errorf("%v: %v", uri, err)
+		}
 	}
-	content, err := ioutil.ReadFile(fname)
-	if err != nil {
-		return nil, fmt.Errorf("%v: %v", uri, err)
-	}
-	f := c.fset.AddFile(fname, -1, len(content))
-	f.SetLinesForContent(content)
-	m := protocol.NewColumnMapper(uri, c.fset, f, content)
-	p := &protocol.DidOpenTextDocumentParams{}
-	p.TextDocument.URI = string(uri)
-	p.TextDocument.Text = string(content)
-	if err := c.server.DidOpen(ctx, p); err != nil {
-		return nil, fmt.Errorf("%v: %v", uri, err)
-	}
-	return m, nil
+	return file
 }
diff --git a/internal/lsp/cmd/definition.go b/internal/lsp/cmd/definition.go
index 4e9b97e..d093ab7 100644
--- a/internal/lsp/cmd/definition.go
+++ b/internal/lsp/cmd/definition.go
@@ -59,17 +59,16 @@
 	if len(args) != 1 {
 		return tool.CommandLineErrorf("definition expects 1 argument")
 	}
-	client := &baseClient{}
-	server, err := d.query.app.connect(ctx, client)
+	conn, err := d.query.app.connect(ctx)
 	if err != nil {
 		return err
 	}
 	from := span.Parse(args[0])
-	m, err := client.AddFile(ctx, from.URI())
-	if err != nil {
-		return err
+	file := conn.AddFile(ctx, from.URI())
+	if file.err != nil {
+		return file.err
 	}
-	loc, err := m.Location(from)
+	loc, err := file.mapper.Location(from)
 	if err != nil {
 		return err
 	}
@@ -77,7 +76,7 @@
 		TextDocument: protocol.TextDocumentIdentifier{URI: loc.URI},
 		Position:     loc.Range.Start,
 	}
-	locs, err := server.Definition(ctx, &p)
+	locs, err := conn.Definition(ctx, &p)
 	if err != nil {
 		return fmt.Errorf("%v: %v", from, err)
 	}
@@ -85,18 +84,18 @@
 	if len(locs) == 0 {
 		return fmt.Errorf("%v: not an identifier", from)
 	}
-	hover, err := server.Hover(ctx, &p)
+	hover, err := conn.Hover(ctx, &p)
 	if err != nil {
 		return fmt.Errorf("%v: %v", from, err)
 	}
 	if hover == nil {
 		return fmt.Errorf("%v: not an identifier", from)
 	}
-	m, err = client.AddFile(ctx, span.NewURI(locs[0].URI))
-	if err != nil {
-		return fmt.Errorf("%v: %v", from, err)
+	file = conn.AddFile(ctx, span.NewURI(locs[0].URI))
+	if file.err != nil {
+		return fmt.Errorf("%v: %v", from, file.err)
 	}
-	definition, err := m.Span(locs[0])
+	definition, err := file.mapper.Span(locs[0])
 	if err != nil {
 		return fmt.Errorf("%v: %v", from, err)
 	}
diff --git a/internal/lsp/cmd/definition_test.go b/internal/lsp/cmd/definition_test.go
index 0419b16..3e44f9a 100644
--- a/internal/lsp/cmd/definition_test.go
+++ b/internal/lsp/cmd/definition_test.go
@@ -34,9 +34,6 @@
 	if runtime.GOOS == "android" {
 		t.Skip("not all source files are available on android")
 	}
-	if runtime.GOOS != "linux" || isRace {
-		t.Skip("currently uses too much memory, see issue #31611")
-	}
 	dir, err := os.Getwd()
 	if err != nil {
 		t.Errorf("could not get wd: %v", err)
@@ -64,7 +61,7 @@
 			// TODO: support type definition queries
 			continue
 		}
-		args := []string{"query"}
+		args := []string{"-remote=internal", "query"}
 		if d.Flags != "" {
 			args = append(args, strings.Split(d.Flags, " ")...)
 		}
diff --git a/internal/lsp/cmd/format.go b/internal/lsp/cmd/format.go
index 32f6468..0f6542e 100644
--- a/internal/lsp/cmd/format.go
+++ b/internal/lsp/cmd/format.go
@@ -50,20 +50,19 @@
 		// no files, so no results
 		return nil
 	}
-	client := &baseClient{}
 	// now we ready to kick things off
-	server, err := f.app.connect(ctx, client)
+	conn, err := f.app.connect(ctx)
 	if err != nil {
 		return err
 	}
 	for _, arg := range args {
 		spn := span.Parse(arg)
-		m, err := client.AddFile(ctx, spn.URI())
-		if err != nil {
-			return err
+		file := conn.AddFile(ctx, spn.URI())
+		if file.err != nil {
+			return file.err
 		}
 		filename, _ := spn.URI().Filename() // this cannot fail, already checked in AddFile above
-		loc, err := m.Location(spn)
+		loc, err := file.mapper.Location(spn)
 		if err != nil {
 			return err
 		}
@@ -73,16 +72,16 @@
 		p := protocol.DocumentFormattingParams{
 			TextDocument: protocol.TextDocumentIdentifier{URI: loc.URI},
 		}
-		edits, err := server.Formatting(ctx, &p)
+		edits, err := conn.Formatting(ctx, &p)
 		if err != nil {
 			return fmt.Errorf("%v: %v", spn, err)
 		}
-		sedits, err := lsp.FromProtocolEdits(m, edits)
+		sedits, err := lsp.FromProtocolEdits(file.mapper, edits)
 		if err != nil {
 			return fmt.Errorf("%v: %v", spn, err)
 		}
 		ops := source.EditsToDiff(sedits)
-		lines := diff.SplitLines(string(m.Content))
+		lines := diff.SplitLines(string(file.mapper.Content))
 		formatted := strings.Join(diff.ApplyEdits(lines, ops), "")
 		printIt := true
 		if f.List {
diff --git a/internal/lsp/cmd/format_test.go b/internal/lsp/cmd/format_test.go
index eaaf8d5..8e38075 100644
--- a/internal/lsp/cmd/format_test.go
+++ b/internal/lsp/cmd/format_test.go
@@ -44,7 +44,7 @@
 			app := &cmd.Application{}
 			app.Config = r.data.Config
 			got := captureStdOut(t, func() {
-				tool.Main(context.Background(), app, append([]string{"format"}, args...))
+				tool.Main(context.Background(), app, append([]string{"-remote=internal", "format"}, args...))
 			})
 			got = r.normalizePaths(got)
 			// check the first two lines are the expected file header