transport.go

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