transport.go

  1package lsp
  2
  3import (
  4	"bufio"
  5	"context"
  6	"encoding/json"
  7	"fmt"
  8	"io"
  9	"log/slog"
 10	"strings"
 11)
 12
 13// WriteMessage writes an LSP message to the given writer
 14func WriteMessage(w io.Writer, msg *Message, debug bool) error {
 15	data, err := json.Marshal(msg)
 16	if err != nil {
 17		return fmt.Errorf("failed to marshal message: %w", err)
 18	}
 19
 20	if debug {
 21		slog.Debug("Sending message to server", "method", msg.Method, "id", msg.ID)
 22	}
 23
 24	_, err = fmt.Fprintf(w, "Content-Length: %d\r\n\r\n", len(data))
 25	if err != nil {
 26		return fmt.Errorf("failed to write header: %w", err)
 27	}
 28
 29	_, err = w.Write(data)
 30	if err != nil {
 31		return fmt.Errorf("failed to write message: %w", err)
 32	}
 33
 34	return nil
 35}
 36
 37// ReadMessage reads a single LSP message from the given reader
 38func ReadMessage(r *bufio.Reader, debug bool) (*Message, error) {
 39	// Read headers
 40	var contentLength int
 41	for {
 42		line, err := r.ReadString('\n')
 43		if err != nil {
 44			return nil, fmt.Errorf("failed to read header: %w", err)
 45		}
 46		line = strings.TrimSpace(line)
 47
 48		if debug {
 49			slog.Debug("Received header", "line", line)
 50		}
 51
 52		if line == "" {
 53			break // End of headers
 54		}
 55
 56		if strings.HasPrefix(line, "Content-Length: ") {
 57			_, err := fmt.Sscanf(line, "Content-Length: %d", &contentLength)
 58			if err != nil {
 59				return nil, fmt.Errorf("invalid Content-Length: %w", err)
 60			}
 61		}
 62	}
 63
 64	if debug {
 65		slog.Debug("Content-Length", "length", contentLength)
 66	}
 67
 68	// Read content
 69	content := make([]byte, contentLength)
 70	_, err := io.ReadFull(r, content)
 71	if err != nil {
 72		return nil, fmt.Errorf("failed to read content: %w", err)
 73	}
 74
 75	if debug {
 76		slog.Debug("Received content", "content", string(content))
 77	}
 78
 79	// Parse message
 80	var msg Message
 81	if err := json.Unmarshal(content, &msg); err != nil {
 82		return nil, fmt.Errorf("failed to unmarshal message: %w", err)
 83	}
 84
 85	return &msg, nil
 86}
 87
 88// handleMessages reads and dispatches messages in a loop
 89func (c *Client) handleMessages() {
 90	for {
 91		msg, err := ReadMessage(c.stdout, c.debug)
 92		if err != nil {
 93			if c.debug {
 94				slog.Error("Error reading message", "error", err)
 95			}
 96			return
 97		}
 98
 99		// Handle server->client request (has both Method and ID)
100		if msg.Method != "" && msg.ID != 0 {
101			if c.debug {
102				slog.Debug("Received request from server", "method", msg.Method, "id", msg.ID)
103			}
104
105			response := &Message{
106				JSONRPC: "2.0",
107				ID:      msg.ID,
108			}
109
110			// Look up handler for this method
111			c.serverHandlersMu.RLock()
112			handler, ok := c.serverRequestHandlers[msg.Method]
113			c.serverHandlersMu.RUnlock()
114
115			if ok {
116				result, err := handler(msg.Params)
117				if err != nil {
118					response.Error = &ResponseError{
119						Code:    -32603,
120						Message: err.Error(),
121					}
122				} else {
123					rawJSON, err := json.Marshal(result)
124					if err != nil {
125						response.Error = &ResponseError{
126							Code:    -32603,
127							Message: fmt.Sprintf("failed to marshal response: %v", err),
128						}
129					} else {
130						response.Result = rawJSON
131					}
132				}
133			} else {
134				response.Error = &ResponseError{
135					Code:    -32601,
136					Message: fmt.Sprintf("method not found: %s", msg.Method),
137				}
138			}
139
140			// Send response back to server
141			if err := WriteMessage(c.stdin, response, c.debug); err != nil {
142				slog.Error("Error sending response to server", "error", err)
143			}
144
145			continue
146		}
147
148		// Handle notification (has Method but no ID)
149		if msg.Method != "" && msg.ID == 0 {
150			c.notificationMu.RLock()
151			handler, ok := c.notificationHandlers[msg.Method]
152			c.notificationMu.RUnlock()
153
154			if ok {
155				if c.debug {
156					slog.Debug("Handling notification", "method", msg.Method)
157				}
158				go handler(msg.Params)
159			} else if c.debug {
160				slog.Debug("No handler for notification", "method", msg.Method)
161			}
162			continue
163		}
164
165		// Handle response to our request (has ID but no Method)
166		if msg.ID != 0 && msg.Method == "" {
167			c.handlersMu.RLock()
168			ch, ok := c.handlers[msg.ID]
169			c.handlersMu.RUnlock()
170
171			if ok {
172				if c.debug {
173					slog.Debug("Received response for request", "id", msg.ID)
174				}
175				ch <- msg
176				close(ch)
177			} else if c.debug {
178				slog.Debug("No handler for response", "id", msg.ID)
179			}
180		}
181	}
182}
183
184// Call makes a request and waits for the response
185func (c *Client) Call(ctx context.Context, method string, params any, result any) error {
186	id := c.nextID.Add(1)
187
188	if c.debug {
189		slog.Debug("Making call", "method", method, "id", id)
190	}
191
192	msg, err := NewRequest(id, method, params)
193	if err != nil {
194		return fmt.Errorf("failed to create request: %w", err)
195	}
196
197	// Create response channel
198	ch := make(chan *Message, 1)
199	c.handlersMu.Lock()
200	c.handlers[id] = ch
201	c.handlersMu.Unlock()
202
203	defer func() {
204		c.handlersMu.Lock()
205		delete(c.handlers, id)
206		c.handlersMu.Unlock()
207	}()
208
209	// Send request
210	if err := WriteMessage(c.stdin, msg, c.debug); err != nil {
211		return fmt.Errorf("failed to send request: %w", err)
212	}
213
214	if c.debug {
215		slog.Debug("Request sent", "method", method, "id", id)
216	}
217
218	// Wait for response
219	select {
220	case <-ctx.Done():
221		return ctx.Err()
222	case resp := <-ch:
223		if c.debug {
224			slog.Debug("Received response", "id", id)
225		}
226
227		if resp.Error != nil {
228			return fmt.Errorf("request failed: %s (code: %d)", resp.Error.Message, resp.Error.Code)
229		}
230
231		if result != nil {
232			// If result is a json.RawMessage, just copy the raw bytes
233			if rawMsg, ok := result.(*json.RawMessage); ok {
234				*rawMsg = resp.Result
235				return nil
236			}
237			// Otherwise unmarshal into the provided type
238			if err := json.Unmarshal(resp.Result, result); err != nil {
239				return fmt.Errorf("failed to unmarshal result: %w", err)
240			}
241		}
242
243		return nil
244	}
245}
246
247// Notify sends a notification (a request without an ID that doesn't expect a response)
248func (c *Client) Notify(ctx context.Context, method string, params any) error {
249	if c.debug {
250		slog.Debug("Sending notification", "method", method)
251	}
252
253	msg, err := NewNotification(method, params)
254	if err != nil {
255		return fmt.Errorf("failed to create notification: %w", err)
256	}
257
258	if err := WriteMessage(c.stdin, msg, c.debug); err != nil {
259		return fmt.Errorf("failed to send notification: %w", err)
260	}
261
262	return nil
263}
264
265type (
266	NotificationHandler  func(params json.RawMessage)
267	ServerRequestHandler func(params json.RawMessage) (any, error)
268)