Skip to content

fix: server lifecycle context [IDE-1791] #1173

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rrama wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix: server lifecycle context [IDE-1791] #1173

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0ba9645
fix: server lifecycle context
rrama Mar 3, 2026
bc5eeee
refactor: pass absolute config file paths
rrama Mar 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions application/config/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,8 @@ type Config struct {
offline bool
ws types.Workspace
isLSPInitialized bool
serverLifecycleCtx context.Context
serverLifecycleCtxCancelFunc context.CancelFunc
cachedOriginalPath string
userSettingsPath string
lastSetOrganization string // Trimmed raw org value last passed to SetOrganization
Expand Down Expand Up @@ -1556,6 +1558,40 @@ func (c *Config) SetLSPInitialized(initialized bool) {
c.isLSPInitialized = initialized
}

func (c *Config) SetServerLifecycleContext(ctx context.Context, cancel context.CancelFunc) {
if ctx == nil {
panic("server lifecycle context must not be nil")
}
if cancel == nil {
panic("server lifecycle cancel func must not be nil")
}
c.m.Lock()
defer c.m.Unlock()
c.serverLifecycleCtx = ctx
c.serverLifecycleCtxCancelFunc = cancel
}

// ServerLifecycleContext returns the language server lifecycle context to be used by
// background operations. It will be canceled when the language server shuts down.
func (c *Config) ServerLifecycleContext() context.Context {
c.m.RLock()
defer c.m.RUnlock()
if c.serverLifecycleCtx == nil {

@bastiandoetsch bastiandoetsch Jun 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Should Fix] These accessors panic when the lifecycle context was never set, but newConfig doesn't default-initialize it — only server.Start() and the test helpers do. Production consumers reached deep in background paths (internal/sdk/sdk.go, infrastructure/oss/cli_scanner.go:602 scheduleRefreshScan) call ServerLifecycleContext() unconditionally, so any Config built via config.New/NewFromExtension that triggers a scan without going through Start() (an embedder, or a unit test bypassing the testutil helpers) hard-panics the whole server instead of degrading. Flagged by 3 of 4 reviewers.

Fix: default-initialize the field in newConfig with context.WithCancel(context.Background()) (overwritten by Start), which lets you delete all three panic guards. (Storing a context on the long-lived shared Config singleton is also the containedctx anti-pattern — pragmatic here since Config is the universal handle, but worth a comment.)

— AI review

panic("server lifecycle context is not initialized")
}
return c.serverLifecycleCtx
}

func (c *Config) CancelServerLifecycleContext() {
c.m.RLock()
cancel := c.serverLifecycleCtxCancelFunc
c.m.RUnlock()
if cancel == nil {
panic("server lifecycle cancel func is not initialized")
}
cancel()
}

func (c *Config) EmptyToken() bool {
return !c.NonEmptyToken()
}
Expand Down
11 changes: 6 additions & 5 deletions application/server/codeaction_handlers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,17 +57,18 @@ func GetCodeActionHandler(c *config.Config) TextDocumentCodeActionHandler {

// We share a mutex between all the handler calls to prevent concurrent runs.
var mu = &sync.Mutex{}
// This "field" is shared between the handlers to allow for cancellation of previous handler
_, cancel := context.WithCancel(context.Background())
// This "field" is shared between the handlers to allow for cancellation of previous handler.
// Use a no-op fake cancel func as the initial cancel func, so the first call has something to cancel.
var cancelFunc context.CancelFunc = func() {}
logger := c.Logger().With().Str("method", "CodeActionHandler").Logger()

return func(paramCtx context.Context, params types.CodeActionParams) ([]types.LSPCodeAction, error) {
// We want to avoid concurrent runs of this handler to prevent race condition.
var ctx context.Context
mu.Lock()
cancel()
ctx, cancel = context.WithCancel(paramCtx)
defer cancel()
cancelFunc()
ctx, cancelFunc = context.WithCancel(paramCtx)
defer cancelFunc()
mu.Unlock()

// This handler uses debouncing because it is called very often on mouse/caret moves.
Expand Down
4 changes: 2 additions & 2 deletions application/server/configuration_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ func Test_WorkspaceDidChangeConfiguration_Push(t *testing.T) {

func Test_WorkspaceDidChangeConfiguration_Pull(t *testing.T) {
c := testutil.UnitTest(t)
loc, _ := setupCustomServer(t, c, callBackMock)
loc, _ := setupServer(t, c, withCallbackFn(callBackMock))

_, err := loc.Client.Call(t.Context(), "initialize", types.InitializeParams{
Capabilities: types.ClientCapabilities{
Expand Down Expand Up @@ -157,7 +157,7 @@ func callBackMock(_ context.Context, request *jrpc2.Request) (any, error) {

func Test_WorkspaceDidChangeConfiguration_PullNoCapability(t *testing.T) {
c := testutil.UnitTest(t)
loc, jsonRPCRecorder := setupCustomServer(t, c, callBackMock)
loc, jsonRPCRecorder := setupServer(t, c, withCallbackFn(callBackMock))

params := types.DidChangeConfigurationParams{Settings: types.Settings{}}
var updated = true
Expand Down
16 changes: 8 additions & 8 deletions application/server/execute_command_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ import (

func Test_executeWorkspaceScanCommand_shouldStartWorkspaceScanOnCommandReceipt(t *testing.T) {
c := testutil.UnitTest(t)
loc, _ := setupServerWithCustomDI(t, c, false)
loc, _ := setupServer(t, c, withRealDI())

s := &scanner.TestScanner{}
c.Workspace().AddFolder(workspace.NewFolder(c, "dummy", "dummy", s, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))
Expand All @@ -59,7 +59,7 @@ func Test_executeWorkspaceScanCommand_shouldStartWorkspaceScanOnCommandReceipt(t

func Test_executeWorkspaceFolderScanCommand_shouldStartFolderScanOnCommandReceipt(t *testing.T) {
c := testutil.UnitTest(t)
loc, _ := setupServerWithCustomDI(t, c, false)
loc, _ := setupServer(t, c, withRealDI())

s := &scanner.TestScanner{}
c.Workspace().AddFolder(workspace.NewFolder(c, "dummy", "dummy", s, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))
Expand All @@ -76,7 +76,7 @@ func Test_executeWorkspaceFolderScanCommand_shouldStartFolderScanOnCommandReceip

func Test_executeWorkspaceFolderScanCommand_shouldNotClearOtherFoldersDiagnostics(t *testing.T) {
c := testutil.UnitTest(t)
loc, _ := setupServerWithCustomDI(t, c, false)
loc, _ := setupServer(t, c, withRealDI())

scannerForFolder := scanner.NewTestScanner()
scannerForDontClear := scanner.NewTestScanner()
Expand Down Expand Up @@ -109,7 +109,7 @@ func Test_executeWorkspaceFolderScanCommand_shouldNotClearOtherFoldersDiagnostic

func Test_executeWorkspaceScanCommand_shouldAskForTrust(t *testing.T) {
c := testutil.UnitTest(t)
loc, jsonRPCRecorder := setupServerWithCustomDI(t, c, false)
loc, jsonRPCRecorder := setupServer(t, c, withRealDI())

s := &scanner.TestScanner{}
c.Workspace().AddFolder(workspace.NewFolder(c, "dummy", "dummy", s, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))
Expand All @@ -128,7 +128,7 @@ func Test_executeWorkspaceScanCommand_shouldAskForTrust(t *testing.T) {

func Test_executeWorkspaceScanCommand_shouldAcceptScanSourceParam(t *testing.T) {
c := testutil.UnitTest(t)
loc, jsonRPCRecorder := setupServerWithCustomDI(t, c, false)
loc, jsonRPCRecorder := setupServer(t, c, withRealDI())

s := &scanner.TestScanner{}
c.Workspace().AddFolder(workspace.NewFolder(c, "dummy", "dummy", s, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))
Expand Down Expand Up @@ -232,7 +232,7 @@ func Test_TrustWorkspaceFolders(t *testing.T) {

t.Run("Doesn't mutate trusted folders, if trusted folders disabled", func(t *testing.T) {
c := testutil.UnitTest(t)
loc, _ := setupServerWithCustomDI(t, c, false)
loc, _ := setupServer(t, c, withRealDI())

c.Workspace().AddFolder(workspace.NewFolder(c, folderPath1, "dummy", nil, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))

Expand All @@ -247,7 +247,7 @@ func Test_TrustWorkspaceFolders(t *testing.T) {

t.Run("Updates trusted workspace folders", func(t *testing.T) {
c := testutil.UnitTest(t)
loc, _ := setupServerWithCustomDI(t, c, false)
loc, _ := setupServer(t, c, withRealDI())

c.Workspace().AddFolder(workspace.NewFolder(c, folderPath1, "dummy", nil, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))
c.Workspace().AddFolder(workspace.NewFolder(c, folderPath2, "dummy", nil, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))
Expand All @@ -266,7 +266,7 @@ func Test_TrustWorkspaceFolders(t *testing.T) {

t.Run("Existing trusted workspace folders are not removed", func(t *testing.T) {
c := testutil.UnitTest(t)
loc, _ := setupServerWithCustomDI(t, c, false)
loc, _ := setupServer(t, c, withRealDI())

c.Workspace().AddFolder(workspace.NewFolder(c, folderPath1, "dummy", nil, di.HoverService(), di.ScanNotifier(), di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver()))
c.SetTrustedFolderFeatureEnabled(true)
Expand Down
4 changes: 2 additions & 2 deletions application/server/notification_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,11 +264,11 @@ func TestShowMessageRequest(t *testing.T) {
t.Run("should execute a command when action item is selected", func(t *testing.T) {
c := testutil.UnitTest(t)
selectedAction := "Open browser"
loc, _ := setupCustomServer(t, c, func(_ context.Context, _ *jrpc2.Request) (any, error) {
loc, _ := setupServer(t, c, withCallbackFn(func(_ context.Context, _ *jrpc2.Request) (any, error) {
return types.MessageActionItem{
Title: selectedAction,
}, nil
})
}))
_, err := loc.Client.Call(t.Context(), "initialize", nil)
if err != nil {
t.Fatal(err)
Expand Down
65 changes: 41 additions & 24 deletions application/server/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,14 @@ import (
func Start(c *config.Config) {
var srv *jrpc2.Server

// Create a context that can be used for background operations that will be canceled when the language server shuts down.
lifecycleCtx, lifecycleCancel := context.WithCancel(context.Background())
c.SetServerLifecycleContext(lifecycleCtx, lifecycleCancel)
// Safety net: if the server exits without processing shutdown or exit,
// still cancel lifecycle-bound background operations. This `defer` will only
// process after `srv.WaitStatus()` returns, which is only after the JRPC server exits.
defer lifecycleCancel()

handlers := handler.Map{}
srv = jrpc2.NewServer(handlers, &jrpc2.ServerOptions{
Logger: func(text string) {
Expand Down Expand Up @@ -97,16 +105,16 @@ func initHandlers(srv *jrpc2.Server, handlers handler.Map, c *config.Config) {
handlers["textDocument/didChange"] = textDocumentDidChangeHandler()
handlers["textDocument/didClose"] = noOpHandler()
handlers[textDocumentDidOpenOperation] = textDocumentDidOpenHandler(c)
handlers[textDocumentDidSaveOperation] = textDocumentDidSaveHandler()
handlers[textDocumentDidSaveOperation] = textDocumentDidSaveHandler(c)
handlers["textDocument/hover"] = textDocumentHover()
handlers["textDocument/codeAction"] = textDocumentCodeActionHandler()
handlers["textDocument/codeLens"] = codeLensHandler()
handlers["textDocument/inlineValue"] = textDocumentInlineValueHandler()
handlers["textDocument/willSave"] = noOpHandler()
handlers["textDocument/willSaveWaitUntil"] = noOpHandler()
handlers["codeAction/resolve"] = codeActionResolveHandler(srv)
handlers["shutdown"] = shutdown()
handlers["exit"] = exit(srv)
handlers["shutdown"] = shutdown(c)
handlers["exit"] = exit(srv, c)
handlers["workspace/didChangeWorkspaceFolders"] = workspaceDidChangeWorkspaceFoldersHandler(c, srv)
handlers["workspace/willDeleteFiles"] = workspaceWillDeleteFilesHandler(c)
handlers["workspace/didChangeConfiguration"] = workspaceDidChangeConfiguration(c, srv)
Expand Down Expand Up @@ -181,23 +189,24 @@ func codeLensHandler() jrpc2.Handler {

func workspaceDidChangeWorkspaceFoldersHandler(c *config.Config, srv *jrpc2.Server) jrpc2.Handler {
return handler.New(func(ctx context.Context, params types.DidChangeWorkspaceFoldersParams) (any, error) {
// The context provided by the JSON-RPC server is canceled once a new message is being processed,
// so we don't want to propagate it to functions that start background operations
bgCtx := context.Background()
logger := c.Logger().With().Str("method", "WorkspaceDidChangeWorkspaceFoldersHandler").Logger()

logger.Info().Msg("RECEIVING")
defer logger.Info().Msg("SENDING")

// The context provided by the JSON-RPC server is canceled once a new message is being processed,
// so we don't want to propagate it to functions that start background operations,
// instead we'll use a background context tied to the language server's lifecycle.
lifecycleCtx := c.ServerLifecycleContext()
changedFolders := c.Workspace().ChangeWorkspaceFolders(params)

if di.AuthenticationService().IsAuthenticated() {
di.LdxSyncService().RefreshConfigFromLdxSync(bgCtx, c, changedFolders, di.Notifier())
di.LdxSyncService().RefreshConfigFromLdxSync(lifecycleCtx, c, changedFolders, di.Notifier())
}

command.HandleFolders(c, bgCtx, srv, di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver())
command.HandleFolders(c, lifecycleCtx, srv, di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver())
for _, f := range changedFolders {
if f.IsAutoScanEnabled() {
go f.ScanFolder(ctx)
go f.ScanFolder(lifecycleCtx)
}
}
return nil, nil
Expand Down Expand Up @@ -443,7 +452,12 @@ func initializedHandler(c *config.Config, srv *jrpc2.Server) handler.Func {
logger.Error().Err(err).Msg("Scan initialization error, canceling scan")
return nil, nil
}
command.HandleFolders(c, context.Background(), srv, di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver())

// The context provided by the JSON-RPC server is canceled once a new message is being processed,
// so we don't want to propagate it to functions that start background operations,
// instead we'll use a background context tied to the language server's lifecycle.
lifecycleCtx := c.ServerLifecycleContext()
command.HandleFolders(c, lifecycleCtx, srv, di.Notifier(), di.ScanPersister(), di.ScanStateAggregator(), di.FeatureFlagService(), di.ConfigResolver())

// Check once for expired cache in same thread before triggering a scan.
// Start a periodic go routine to check for the expired cache afterwards
Expand All @@ -453,7 +467,7 @@ func initializedHandler(c *config.Config, srv *jrpc2.Server) handler.Func {
autoScanEnabled := c.IsAutoScanEnabled()
if autoScanEnabled {
logger.Info().Msg("triggering workspace scan after successful initialization")
c.Workspace().ScanWorkspace(context.Background())
c.Workspace().ScanWorkspace(lifecycleCtx)
} else {
msg := fmt.Sprintf(
"No automatic workspace scan on initialization: autoScanEnabled=%v",
Expand Down Expand Up @@ -630,12 +644,12 @@ func monitorClientProcess(pid int) time.Duration {
return time.Since(start)
}

func shutdown() jrpc2.Handler {
return handler.New(func(ctx context.Context) (any, error) {
c := config.CurrentConfig()
func shutdown(c *config.Config) jrpc2.Handler {
return handler.New(func(_ context.Context) (any, error) {
logger := c.Logger().With().Str("method", "Shutdown").Logger()
logger.Info().Msg("ENTERING")
defer logger.Info().Msg("RETURNING")
c.CancelServerLifecycleContext()
di.ErrorReporter().FlushErrorReporting()

disposeProgressListener()
Expand All @@ -645,11 +659,13 @@ func shutdown() jrpc2.Handler {
})
}

func exit(srv *jrpc2.Server) jrpc2.Handler {
func exit(srv *jrpc2.Server, c *config.Config) jrpc2.Handler {
return handler.New(func(_ context.Context) (any, error) {
c := config.CurrentConfig()
logger := c.Logger().With().Str("method", "Exit").Logger()
logger.Info().Msg("ENTERING")
// Canceling the server's lifecycle context should have been covered by shutdown,
// but we'll do it again, just in case it was skipped.
c.CancelServerLifecycleContext()
logger.Info().Msg("Flushing error reporting...")
di.ErrorReporter().FlushErrorReporting()
logger.Info().Msg("Stopping server...")
Expand Down Expand Up @@ -704,15 +720,16 @@ func textDocumentDidOpenHandler(c *config.Config) jrpc2.Handler {
})
}

func textDocumentDidSaveHandler() jrpc2.Handler {
func textDocumentDidSaveHandler(c *config.Config) jrpc2.Handler {
return handler.New(func(_ context.Context, params sglsp.DidSaveTextDocumentParams) (any, error) {
// The context provided by the JSON-RPC server is canceled once a new message is being processed,
// so we don't want to propagate it to functions that start background operations
bgCtx := context.Background()
c := config.CurrentConfig()
logger := c.Logger().With().Str("method", "TextDocumentDidSaveHandler").Logger()
logger.Debug().Interface("params", params).Msg("Receiving")

// The context provided by the JSON-RPC server is canceled once a new message is being processed,
// so we don't want to propagate it to functions that start background operations,
// instead we'll use a background context tied to the language server's lifecycle.
lifecycleCtx := c.ServerLifecycleContext()

di.FileWatcher().SetFileAsSaved(params.TextDocument.URI)
filePath := uri.PathFromUri(params.TextDocument.URI)

Expand All @@ -728,12 +745,12 @@ func textDocumentDidSaveHandler() jrpc2.Handler {
}

if folder.IsAutoScanEnabled() && uri.IsDotSnykFile(params.TextDocument.URI) {
go folder.ScanFolder(bgCtx)
go folder.ScanFolder(lifecycleCtx)
return nil, nil
}

if folder.IsAutoScanEnabled() {
go folder.ScanFile(bgCtx, filePath)
go folder.ScanFile(lifecycleCtx, filePath)
} else {
logger.Warn().Msg("Not scanning, auto-scan is disabled")
}
Expand Down
Loading
Loading