chore: pass context to lsp state

Kujtim Hoxha created

Change summary

internal/app/lsp.go        | 12 ++++++------
internal/app/lsp_events.go |  8 ++++----
internal/lsp/client.go     |  8 ++++----
internal/lsp/handlers.go   |  4 ++--
4 files changed, 16 insertions(+), 16 deletions(-)

Detailed changes

internal/app/lsp.go 🔗

@@ -28,18 +28,18 @@ func (app *App) createAndStartLSPClient(ctx context.Context, name string, config
 	// Check if any root markers exist in the working directory (config now has defaults)
 	if !lsp.HasRootMarkers(app.config.WorkingDir(), config.RootMarkers) {
 		slog.Debug("Skipping LSP client: no root markers found", "name", name, "rootMarkers", config.RootMarkers)
-		updateLSPState(name, lsp.StateDisabled, nil, nil, 0)
+		updateLSPState(ctx, name, lsp.StateDisabled, nil, nil, 0)
 		return
 	}
 
 	// Update state to starting
-	updateLSPState(name, lsp.StateStarting, nil, nil, 0)
+	updateLSPState(ctx, name, lsp.StateStarting, nil, nil, 0)
 
 	// Create LSP client.
 	lspClient, err := lsp.New(ctx, name, config, app.config.Resolver())
 	if err != nil {
 		slog.Error("Failed to create LSP client for", "name", name, "error", err)
-		updateLSPState(name, lsp.StateError, err, nil, 0)
+		updateLSPState(ctx, name, lsp.StateError, err, nil, 0)
 		return
 	}
 
@@ -54,7 +54,7 @@ func (app *App) createAndStartLSPClient(ctx context.Context, name string, config
 	_, err = lspClient.Initialize(initCtx, app.config.WorkingDir())
 	if err != nil {
 		slog.Error("LSP client initialization failed", "name", name, "error", err)
-		updateLSPState(name, lsp.StateError, err, lspClient, 0)
+		updateLSPState(ctx, name, lsp.StateError, err, lspClient, 0)
 		lspClient.Close(ctx)
 		return
 	}
@@ -65,12 +65,12 @@ func (app *App) createAndStartLSPClient(ctx context.Context, name string, config
 		// Server never reached a ready state, but let's continue anyway, as
 		// some functionality might still work.
 		lspClient.SetServerState(lsp.StateError)
-		updateLSPState(name, lsp.StateError, err, lspClient, 0)
+		updateLSPState(ctx, name, lsp.StateError, err, lspClient, 0)
 	} else {
 		// Server reached a ready state successfully.
 		slog.Debug("LSP server is ready", "name", name)
 		lspClient.SetServerState(lsp.StateReady)
-		updateLSPState(name, lsp.StateReady, nil, lspClient, 0)
+		updateLSPState(ctx, name, lsp.StateReady, nil, lspClient, 0)
 	}
 
 	slog.Info("LSP client initialized", "name", name)

internal/app/lsp_events.go 🔗

@@ -57,7 +57,7 @@ func GetLSPState(name string) (LSPClientInfo, bool) {
 }
 
 // updateLSPState updates the state of an LSP client and publishes an event
-func updateLSPState(name string, state lsp.ServerState, err error, client *lsp.Client, diagnosticCount int) {
+func updateLSPState(ctx context.Context, name string, state lsp.ServerState, err error, client *lsp.Client, diagnosticCount int) {
 	info := LSPClientInfo{
 		Name:            name,
 		State:           state,
@@ -71,7 +71,7 @@ func updateLSPState(name string, state lsp.ServerState, err error, client *lsp.C
 	lspStates.Set(name, info)
 
 	// Publish state change event
-	lspBroker.Publish(context.Background(), pubsub.UpdatedEvent, LSPEvent{
+	lspBroker.Publish(ctx, pubsub.UpdatedEvent, LSPEvent{
 		Type:            LSPEventStateChanged,
 		Name:            name,
 		State:           state,
@@ -81,13 +81,13 @@ func updateLSPState(name string, state lsp.ServerState, err error, client *lsp.C
 }
 
 // updateLSPDiagnostics updates the diagnostic count for an LSP client and publishes an event
-func updateLSPDiagnostics(name string, diagnosticCount int) {
+func updateLSPDiagnostics(ctx context.Context, name string, diagnosticCount int) {
 	if info, exists := lspStates.Get(name); exists {
 		info.DiagnosticCount = diagnosticCount
 		lspStates.Set(name, info)
 
 		// Publish diagnostics change event
-		lspBroker.Publish(context.Background(), pubsub.UpdatedEvent, LSPEvent{
+		lspBroker.Publish(ctx, pubsub.UpdatedEvent, LSPEvent{
 			Type:            LSPEventDiagnosticsChanged,
 			Name:            name,
 			State:           info.State,

internal/lsp/client.go 🔗

@@ -41,7 +41,7 @@ type Client struct {
 	config config.LSPConfig
 
 	// Diagnostic change callback
-	onDiagnosticsChanged func(name string, count int)
+	onDiagnosticsChanged func(ctx context.Context, name string, count int)
 
 	// Diagnostic cache
 	diagnostics *csync.VersionedMap[protocol.DocumentURI, []protocol.Diagnostic]
@@ -144,8 +144,8 @@ func (c *Client) Initialize(ctx context.Context, workspaceDir string) (*protocol
 	c.RegisterServerRequestHandler("workspace/configuration", HandleWorkspaceConfiguration)
 	c.RegisterServerRequestHandler("client/registerCapability", HandleRegisterCapability)
 	c.RegisterNotificationHandler("window/showMessage", HandleServerMessage)
-	c.RegisterNotificationHandler("textDocument/publishDiagnostics", func(_ context.Context, _ string, params json.RawMessage) {
-		HandleDiagnostics(c, params)
+	c.RegisterNotificationHandler("textDocument/publishDiagnostics", func(ctx context.Context, _ string, params json.RawMessage) {
+		HandleDiagnostics(ctx, c, params)
 	})
 
 	return result, nil
@@ -192,7 +192,7 @@ func (c *Client) GetName() string {
 }
 
 // SetDiagnosticsCallback sets the callback function for diagnostic changes
-func (c *Client) SetDiagnosticsCallback(callback func(name string, count int)) {
+func (c *Client) SetDiagnosticsCallback(callback func(ctx context.Context, name string, count int)) {
 	c.onDiagnosticsChanged = callback
 }
 

internal/lsp/handlers.go 🔗

@@ -104,7 +104,7 @@ func HandleServerMessage(_ context.Context, method string, params json.RawMessag
 }
 
 // HandleDiagnostics handles diagnostic notifications from the LSP server
-func HandleDiagnostics(client *Client, params json.RawMessage) {
+func HandleDiagnostics(ctx context.Context, client *Client, params json.RawMessage) {
 	var diagParams protocol.PublishDiagnosticsParams
 	if err := json.Unmarshal(params, &diagParams); err != nil {
 		slog.Error("Error unmarshaling diagnostics params", "error", err)
@@ -121,6 +121,6 @@ func HandleDiagnostics(client *Client, params json.RawMessage) {
 
 	// Trigger callback if set
 	if client.onDiagnosticsChanged != nil {
-		client.onDiagnosticsChanged(client.name, totalCount)
+		client.onDiagnosticsChanged(ctx, client.name, totalCount)
 	}
 }