From 255e98f6bb710a3ed33bf23593a2633903bc3968 Mon Sep 17 00:00:00 2001 From: Andrew Branch Date: Wed, 11 Feb 2026 17:31:47 -0800 Subject: [PATCH 1/4] Acquire snapshots synchronously while blocking other requests/notifications --- internal/lsp/server.go | 193 ++++++++++++++++++++++------------------- 1 file changed, 102 insertions(+), 91 deletions(-) diff --git a/internal/lsp/server.go b/internal/lsp/server.go index 28abbe6a9e..edd42cfbb0 100644 --- a/internal/lsp/server.go +++ b/internal/lsp/server.go @@ -415,32 +415,34 @@ func (s *Server) dispatchLoop(ctx context.Context) error { s.pendingClientRequestsMu.Unlock() } - handle := func() { - if err := s.handleRequestOrNotification(requestCtx, req); err != nil { - if errors.Is(err, context.Canceled) { - if err := s.sendError(req.ID, lsproto.ErrorCodeRequestCancelled); err != nil { - lspExit(err) - } - } else if errors.Is(err, io.EOF) { - lspExit(nil) - } else { - if err := s.sendError(req.ID, err); err != nil { - lspExit(err) - } + handleError := func(err error) { + if errors.Is(err, context.Canceled) { + if err := s.sendError(req.ID, lsproto.ErrorCodeRequestCancelled); err != nil { + lspExit(err) + } + } else if errors.Is(err, io.EOF) { + lspExit(nil) + } else { + if err := s.sendError(req.ID, err); err != nil { + lspExit(err) } } + } - if req.ID != nil { - s.pendingClientRequestsMu.Lock() - delete(s.pendingClientRequests, *req.ID) - s.pendingClientRequestsMu.Unlock() - } + if doAsyncWork, err := s.handleRequestOrNotification(requestCtx, req); err != nil { + handleError(err) + } else if doAsyncWork != nil { + go func() { + if lsError := doAsyncWork(); lsError != nil { + handleError(lsError) + } + }() } - if isBlockingMethod(req.Method) { - handle() - } else { - go handle() + if req.ID != nil { + s.pendingClientRequestsMu.Lock() + delete(s.pendingClientRequests, *req.ID) + s.pendingClientRequestsMu.Unlock() } } } @@ -532,27 +534,46 @@ func (s *Server) send(msg *lsproto.Message) error { } } -func (s *Server) handleRequestOrNotification(ctx context.Context, req *lsproto.RequestMessage) error { +// handleRequestOrNotification looks up the handler for the given request or notification, executes its synchronous work +// and returns any asynchronous work as a function to be executed by the caller. +func (s *Server) handleRequestOrNotification(ctx context.Context, req *lsproto.RequestMessage) (func() error, error) { ctx = lsproto.WithClientCapabilities(ctx, &s.clientCapabilities) if handler := handlers()[req.Method]; handler != nil { start := time.Now() - err := handler(s, ctx, req) + doAsyncWork, err := handler(s, ctx, req) idStr := "" if req.ID != nil { idStr = " (" + req.ID.String() + ")" } + if err != nil { + s.logger.Error("error handling method '", req.Method, "'", idStr, ": ", err) + return nil, err + } + if doAsyncWork != nil { + return func() error { + if ctx.Err() != nil { + return ctx.Err() + } + asyncWorkErr := doAsyncWork() + s.logger.Info(core.IfElse(asyncWorkErr != nil, "error handling method '", "handled method '"), req.Method, "'", idStr, " in ", time.Since(start)) + return asyncWorkErr + }, nil + } s.logger.Info("handled method '", req.Method, "'", idStr, " in ", time.Since(start)) - return err + return nil, nil } s.logger.Warn("unknown method '", req.Method, "'") if req.ID != nil { - return s.sendError(req.ID, lsproto.ErrorCodeInvalidRequest) + return nil, s.sendError(req.ID, lsproto.ErrorCodeInvalidRequest) } - return nil + return nil, nil } -type handlerMap map[lsproto.Method]func(*Server, context.Context, *lsproto.RequestMessage) error +// handlerMap maps LSP method to a handler function. The handler function executes any work that must be done synchronously +// before other requests/notifications can be processed, and returns any additional work as a function to be executed +// asynchronously after the synchronous work is complete. +type handlerMap map[lsproto.Method]func(*Server, context.Context, *lsproto.RequestMessage) (func() error, error) var handlers = sync.OnceValue(func() handlerMap { handlers := make(handlerMap) @@ -615,9 +636,9 @@ var handlers = sync.OnceValue(func() handlerMap { }) func registerNotificationHandler[Req any](handlers handlerMap, info lsproto.NotificationInfo[Req], fn func(*Server, context.Context, Req) error) { - handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) error { + handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) (func() error, error) { if s.session == nil && req.Method != lsproto.MethodInitialized { - return lsproto.ErrorCodeServerNotInitialized + return nil, lsproto.ErrorCodeServerNotInitialized } var params Req @@ -626,9 +647,9 @@ func registerNotificationHandler[Req any](handlers handlerMap, info lsproto.Noti params = req.Params.(Req) } if err := fn(s, ctx, params); err != nil { - return err + return nil, err } - return ctx.Err() + return nil, ctx.Err() } } @@ -637,9 +658,9 @@ func registerRequestHandler[Req, Resp any]( info lsproto.RequestInfo[Req, Resp], fn func(*Server, context.Context, Req, *lsproto.RequestMessage) (Resp, error), ) { - handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) error { + handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) (func() error, error) { if s.session == nil && req.Method != lsproto.MethodInitialize { - return lsproto.ErrorCodeServerNotInitialized + return nil, lsproto.ErrorCodeServerNotInitialized } var params Req @@ -649,17 +670,17 @@ func registerRequestHandler[Req, Resp any]( } resp, err := fn(s, ctx, params, req) if err != nil { - return err + return nil, err } if ctx.Err() != nil { - return ctx.Err() + return nil, ctx.Err() } - return s.sendResult(req.ID, resp) + return nil, s.sendResult(req.ID, resp) } } func registerLanguageServiceDocumentRequestHandler[Req lsproto.HasTextDocumentURI, Resp any](handlers handlerMap, info lsproto.RequestInfo[Req, Resp], fn func(*Server, context.Context, *ls.LanguageService, Req) (Resp, error)) { - handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) error { + handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) (func() error, error) { var params Req // Ignore empty params. if req.Params != nil { @@ -667,22 +688,24 @@ func registerLanguageServiceDocumentRequestHandler[Req lsproto.HasTextDocumentUR } ls, err := s.session.GetLanguageService(ctx, params.TextDocumentURI()) if err != nil { - return err + return nil, err } - defer s.recover(ctx, req) - resp, err := fn(s, ctx, ls, params) - if err != nil { - return err - } - if ctx.Err() != nil { - return ctx.Err() - } - return s.sendResult(req.ID, resp) + return func() error { + defer s.recover(ctx, req) + resp, lsErr := fn(s, ctx, ls, params) + if lsErr != nil { + return lsErr + } + if ctx.Err() != nil { + return ctx.Err() + } + return s.sendResult(req.ID, resp) + }, nil } } func registerLanguageServiceWithAutoImportsRequestHandler[Req lsproto.HasTextDocumentURI, Resp any](handlers handlerMap, info lsproto.RequestInfo[Req, Resp], fn func(*Server, context.Context, *ls.LanguageService, Req) (Resp, error)) { - handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) error { + handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) (func() error, error) { var params Req // Ignore empty params. if req.Params != nil { @@ -690,30 +713,32 @@ func registerLanguageServiceWithAutoImportsRequestHandler[Req lsproto.HasTextDoc } languageService, err := s.session.GetLanguageService(ctx, params.TextDocumentURI()) if err != nil { - return err + return nil, err } - defer s.recover(ctx, req) - resp, err := fn(s, ctx, languageService, params) - if errors.Is(err, ls.ErrNeedsAutoImports) { - languageService, err = s.session.GetLanguageServiceWithAutoImports(ctx, params.TextDocumentURI()) - if err != nil { - return err + return func() error { + defer s.recover(ctx, req) + resp, lsErr := fn(s, ctx, languageService, params) + if errors.Is(lsErr, ls.ErrNeedsAutoImports) { + languageService, lsErr = s.session.GetLanguageServiceWithAutoImports(ctx, params.TextDocumentURI()) + if lsErr != nil { + return lsErr + } + if ctx.Err() != nil { + return ctx.Err() + } + resp, lsErr = fn(s, ctx, languageService, params) + if errors.Is(lsErr, ls.ErrNeedsAutoImports) { + panic(info.Method + " returned ErrNeedsAutoImports even after enabling auto imports") + } + } + if lsErr != nil { + return lsErr } if ctx.Err() != nil { return ctx.Err() } - resp, err = fn(s, ctx, languageService, params) - if errors.Is(err, ls.ErrNeedsAutoImports) { - panic(info.Method + " returned ErrNeedsAutoImports even after enabling auto imports") - } - } - if err != nil { - return err - } - if ctx.Err() != nil { - return ctx.Err() - } - return s.sendResult(req.ID, resp) + return s.sendResult(req.ID, resp) + }, nil } } @@ -722,7 +747,7 @@ func registerMultiProjectReferenceRequestHandler[Req lsproto.HasTextDocumentPosi info lsproto.RequestInfo[Req, Resp], fn func(*ls.LanguageService, context.Context, Req, ls.CrossProjectOrchestrator) (Resp, error), ) { - handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) error { + handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) (func() error, error) { var params Req // Ignore empty params. if req.Params != nil { @@ -731,14 +756,16 @@ func registerMultiProjectReferenceRequestHandler[Req lsproto.HasTextDocumentPosi // !!! sheetal: multiple projects that contain the file through symlinks defaultLs, orchestrator, err := s.getLanguageServiceAndCrossProjectOrchestrator(ctx, params.TextDocumentURI(), req) if err != nil { - return err - } - defer s.recover(ctx, req) - resp, err := fn(defaultLs, ctx, params, orchestrator) - if err != nil { - return err + return nil, err } - return s.sendResult(req.ID, resp) + return func() error { + defer s.recover(ctx, req) + resp, lsErr := fn(defaultLs, ctx, params, orchestrator) + if lsErr != nil { + return lsErr + } + return s.sendResult(req.ID, resp) + }, nil } } @@ -1343,22 +1370,6 @@ func (s *Server) NpmInstall(cwd string, args []string) ([]byte, error) { return s.npmInstall(cwd, args) } -func isBlockingMethod(method lsproto.Method) bool { - switch method { - case lsproto.MethodInitialize, - lsproto.MethodInitialized, - lsproto.MethodTextDocumentDidOpen, - lsproto.MethodTextDocumentDidChange, - lsproto.MethodTextDocumentDidSave, - lsproto.MethodTextDocumentDidClose, - lsproto.MethodWorkspaceDidChangeWatchedFiles, - lsproto.MethodWorkspaceDidChangeConfiguration, - lsproto.MethodWorkspaceConfiguration: - return true - } - return false -} - func ptrTo[T any](v T) *T { return &v } From bbb3f2b6817b74a94f66cfa185b81e37a7d45375 Mon Sep 17 00:00:00 2001 From: Andrew Branch Date: Thu, 12 Feb 2026 14:06:04 -0800 Subject: [PATCH 2/4] Move pendingClientRequests deletion to defer after async work Implements suggestion from https://github.com/microsoft/typescript-go/pull/2765#discussion_r2800279862 This ensures that the request is tracked in pendingClientRequests until both the sync and async portions of the work complete, preventing the entry from being deleted while async work is still running. --- internal/lsp/server.go | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/internal/lsp/server.go b/internal/lsp/server.go index 6803ab2de9..2690cb0bc7 100644 --- a/internal/lsp/server.go +++ b/internal/lsp/server.go @@ -429,21 +429,25 @@ func (s *Server) dispatchLoop(ctx context.Context) error { } } - if doAsyncWork, err := s.handleRequestOrNotification(requestCtx, req); err != nil { - handleError(err) - } else if doAsyncWork != nil { - go func() { - if lsError := doAsyncWork(); lsError != nil { - handleError(lsError) + func() { + defer func() { + if req.ID != nil { + s.pendingClientRequestsMu.Lock() + defer s.pendingClientRequestsMu.Unlock() + delete(s.pendingClientRequests, *req.ID) } }() - } - if req.ID != nil { - s.pendingClientRequestsMu.Lock() - delete(s.pendingClientRequests, *req.ID) - s.pendingClientRequestsMu.Unlock() - } + if doAsyncWork, err := s.handleRequestOrNotification(requestCtx, req); err != nil { + handleError(err) + } else if doAsyncWork != nil { + go func() { + if lsError := doAsyncWork(); lsError != nil { + handleError(lsError) + } + }() + } + }() } } } From 11d55c5f553ba1a606fd07bb58c29192f0559b74 Mon Sep 17 00:00:00 2001 From: Andrew Branch Date: Thu, 12 Feb 2026 15:22:15 -0800 Subject: [PATCH 3/4] Fix merge --- internal/lsp/server.go | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/internal/lsp/server.go b/internal/lsp/server.go index 2690cb0bc7..2429743f8d 100644 --- a/internal/lsp/server.go +++ b/internal/lsp/server.go @@ -1374,26 +1374,6 @@ func (s *Server) NpmInstall(cwd string, args []string) ([]byte, error) { return s.npmInstall(cwd, args) } -func isBlockingMethod(method lsproto.Method) bool { - switch method { - case lsproto.MethodInitialize, - lsproto.MethodInitialized, - lsproto.MethodTextDocumentDidOpen, - lsproto.MethodTextDocumentDidChange, - lsproto.MethodTextDocumentDidSave, - lsproto.MethodTextDocumentDidClose, - lsproto.MethodWorkspaceDidChangeWatchedFiles, - lsproto.MethodWorkspaceDidChangeConfiguration, - lsproto.MethodWorkspaceConfiguration: - return true - } - return false -} - -func ptrTo[T any](v T) *T { - return &v -} - // Developer/debugging command handlers func (s *Server) handleRunGC(_ context.Context, _ any, _ *lsproto.RequestMessage) (lsproto.RunGCResponse, error) { From 8e492b5a3609c7def6399a3966760f4369d64fee Mon Sep 17 00:00:00 2001 From: Andrew Branch Date: Thu, 12 Feb 2026 15:59:18 -0800 Subject: [PATCH 4/4] Keep requests cancelable until async work is done --- internal/lsp/server.go | 36 +++++++++++++++++++----------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/internal/lsp/server.go b/internal/lsp/server.go index 2429743f8d..3142094c6b 100644 --- a/internal/lsp/server.go +++ b/internal/lsp/server.go @@ -429,25 +429,27 @@ func (s *Server) dispatchLoop(ctx context.Context) error { } } - func() { - defer func() { - if req.ID != nil { - s.pendingClientRequestsMu.Lock() - defer s.pendingClientRequestsMu.Unlock() - delete(s.pendingClientRequests, *req.ID) + removeRequest := func() { + if req.ID != nil { + s.pendingClientRequestsMu.Lock() + defer s.pendingClientRequestsMu.Unlock() + delete(s.pendingClientRequests, *req.ID) + } + } + + if doAsyncWork, err := s.handleRequestOrNotification(requestCtx, req); err != nil { + handleError(err) + removeRequest() + } else if doAsyncWork != nil { + go func() { + if lsError := doAsyncWork(); lsError != nil { + handleError(lsError) } + removeRequest() }() - - if doAsyncWork, err := s.handleRequestOrNotification(requestCtx, req); err != nil { - handleError(err) - } else if doAsyncWork != nil { - go func() { - if lsError := doAsyncWork(); lsError != nil { - handleError(lsError) - } - }() - } - }() + } else { + removeRequest() + } } } }