jsonrpc2.go

package jsonrpc2

import (
	"bufio"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"reflect"
	"sync"
	"sync/atomic"
)

// Error represents a JSON-RPC error response
type Error struct {
	Code    int            `json:"code"`
	Message string         `json:"message"`
	Data    map[string]any `json:"data,omitempty"`
}

func (e *Error) Error() string {
	return fmt.Sprintf("JSON-RPC Error %d: %s", e.Code, e.Message)
}

// Request represents a JSON-RPC 2.0 request
type Request struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"` // nil for notifications
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params"`
}

func (r *Request) IsCall() bool {
	return len(r.ID) > 0
}

// Response represents a JSON-RPC 2.0 response
type Response struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Result  json.RawMessage `json:"result,omitempty"`
	Error   *Error          `json:"error,omitempty"`
}

// NotificationHandler handles incoming notifications
type NotificationHandler func(method string, params json.RawMessage)

// RequestHandler handles incoming server requests and returns a result or error
type RequestHandler func(params json.RawMessage) (json.RawMessage, *Error)

// Client is a minimal JSON-RPC 2.0 client for stdio transport
type Client struct {
	stdin           io.WriteCloser
	stdout          io.ReadCloser
	mu              sync.Mutex
	pendingRequests map[string]chan *Response
	requestHandlers map[string]RequestHandler
	running         atomic.Bool
	stopChan        chan struct{}
	wg              sync.WaitGroup
	processDone     chan struct{} // closed when the underlying process exits
	processErrorPtr *error        // points to the process error
	processErrorMu  sync.RWMutex  // protects processErrorPtr
	onClose         func()        // called when the read loop exits unexpectedly
}

// NewClient creates a new JSON-RPC client
func NewClient(stdin io.WriteCloser, stdout io.ReadCloser) *Client {
	return &Client{
		stdin:           stdin,
		stdout:          stdout,
		pendingRequests: make(map[string]chan *Response),
		requestHandlers: make(map[string]RequestHandler),
		stopChan:        make(chan struct{}),
	}
}

// SetProcessDone sets a channel that will be closed when the process exits,
// and stores the error pointer that should be returned to pending/future requests.
// The error is read directly from the pointer after the channel closes, avoiding
// a race between an async goroutine and callers checking the error.
func (c *Client) SetProcessDone(done chan struct{}, errPtr *error) {
	c.processDone = done
	c.processErrorMu.Lock()
	c.processErrorPtr = errPtr
	c.processErrorMu.Unlock()
}

// getProcessError returns the process exit error if the process has exited.
// It reads directly from the stored error pointer, which is guaranteed to be
// set before the processDone channel is closed.
func (c *Client) getProcessError() error {
	c.processErrorMu.RLock()
	defer c.processErrorMu.RUnlock()
	if c.processErrorPtr != nil {
		return *c.processErrorPtr
	}
	return nil
}

// Start begins listening for messages in a background goroutine
func (c *Client) Start() {
	c.running.Store(true)
	c.wg.Add(1)
	go c.readLoop()
}

// Stop stops the client and cleans up
func (c *Client) Stop() {
	if !c.running.Load() {
		return
	}
	c.running.Store(false)
	close(c.stopChan)

	// Close stdout to unblock the readLoop
	if c.stdout != nil {
		c.stdout.Close()
	}

	c.wg.Wait()
}

func NotificationHandlerFor[In any](handler func(params In)) RequestHandler {
	return func(params json.RawMessage) (json.RawMessage, *Error) {
		var in In
		// If In is a pointer type, allocate the underlying value and unmarshal into it directly
		var target any = &in
		if t := reflect.TypeFor[In](); t.Kind() == reflect.Pointer {
			in = reflect.New(t.Elem()).Interface().(In)
			target = in
		}
		if err := json.Unmarshal(params, target); err != nil {
			return nil, &Error{
				Code:    -32602,
				Message: fmt.Sprintf("Invalid params: %v", err),
			}
		}
		handler(in)
		return nil, nil
	}
}

// RequestHandlerFor creates a RequestHandler from a typed function
func RequestHandlerFor[In, Out any](handler func(params In) (Out, *Error)) RequestHandler {
	return func(params json.RawMessage) (json.RawMessage, *Error) {
		var in In
		// If In is a pointer type, allocate the underlying value and unmarshal into it directly
		var target any = &in
		if t := reflect.TypeOf(in); t != nil && t.Kind() == reflect.Pointer {
			in = reflect.New(t.Elem()).Interface().(In)
			target = in
		}
		if err := json.Unmarshal(params, target); err != nil {
			return nil, &Error{
				Code:    -32602,
				Message: fmt.Sprintf("Invalid params: %v", err),
			}
		}
		out, errj := handler(in)
		if errj != nil {
			return nil, errj
		}
		outData, err := json.Marshal(out)
		if err != nil {
			return nil, &Error{
				Code:    -32603,
				Message: fmt.Sprintf("Failed to marshal response: %v", err),
			}
		}
		return outData, nil
	}
}

// SetRequestHandler registers a handler for incoming requests from the server
func (c *Client) SetRequestHandler(method string, handler RequestHandler) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if handler == nil {
		delete(c.requestHandlers, method)
		return
	}
	c.requestHandlers[method] = handler
}

// Request sends a JSON-RPC request and waits for the response
func (c *Client) Request(method string, params any) (json.RawMessage, error) {
	requestID := generateUUID()

	// Create response channel
	responseChan := make(chan *Response, 1)
	c.mu.Lock()
	c.pendingRequests[requestID] = responseChan
	c.mu.Unlock()

	// Clean up on exit
	defer func() {
		c.mu.Lock()
		delete(c.pendingRequests, requestID)
		c.mu.Unlock()
	}()

	// Check if process already exited before sending
	if c.processDone != nil {
		select {
		case <-c.processDone:
			if err := c.getProcessError(); err != nil {
				return nil, err
			}
			return nil, fmt.Errorf("process exited unexpectedly")
		default:
			// Process still running, continue
		}
	}

	paramsData, err := json.Marshal(params)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal params: %w", err)
	}

	// Send request
	request := Request{
		JSONRPC: "2.0",
		ID:      json.RawMessage(`"` + requestID + `"`),
		Method:  method,
		Params:  json.RawMessage(paramsData),
	}

	if err := c.sendMessage(request); err != nil {
		return nil, fmt.Errorf("failed to send request: %w", err)
	}

	// Wait for response, also checking for process exit
	if c.processDone != nil {
		select {
		case response := <-responseChan:
			if response.Error != nil {
				return nil, response.Error
			}
			return response.Result, nil
		case <-c.processDone:
			if err := c.getProcessError(); err != nil {
				return nil, err
			}
			return nil, fmt.Errorf("process exited unexpectedly")
		case <-c.stopChan:
			return nil, fmt.Errorf("client stopped")
		}
	}
	select {
	case response := <-responseChan:
		if response.Error != nil {
			return nil, response.Error
		}
		return response.Result, nil
	case <-c.stopChan:
		return nil, fmt.Errorf("client stopped")
	}
}

// Notify sends a JSON-RPC notification (no response expected)
func (c *Client) Notify(method string, params any) error {
	paramsData, err := json.Marshal(params)
	if err != nil {
		return fmt.Errorf("failed to marshal params: %w", err)
	}

	notification := Request{
		JSONRPC: "2.0",
		Method:  method,
		Params:  json.RawMessage(paramsData),
	}
	return c.sendMessage(notification)
}

// sendMessage writes a message to stdin
func (c *Client) sendMessage(message any) error {
	data, err := json.Marshal(message)
	if err != nil {
		return fmt.Errorf("failed to marshal message: %w", err)
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	// Write Content-Length header + message
	header := fmt.Sprintf("Content-Length: %d\r\n\r\n", len(data))
	if _, err := c.stdin.Write([]byte(header)); err != nil {
		return fmt.Errorf("failed to write header: %w", err)
	}
	if _, err := c.stdin.Write(data); err != nil {
		return fmt.Errorf("failed to write message: %w", err)
	}

	return nil
}

// SetOnClose sets a callback invoked when the read loop exits unexpectedly
// (e.g. the underlying connection or process was lost).
func (c *Client) SetOnClose(fn func()) {
	c.onClose = fn
}

// readLoop reads messages from stdout in a background goroutine
func (c *Client) readLoop() {
	defer c.wg.Done()
	defer func() {
		// If still running, the read loop exited unexpectedly (process died or
		// connection dropped). Notify the caller so it can update its state.
		if c.onClose != nil && c.running.Load() {
			c.onClose()
		}
	}()

	reader := bufio.NewReader(c.stdout)

	for c.running.Load() {
		// Read Content-Length header
		var contentLength int
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				// Only log unexpected errors (not EOF or closed pipe during shutdown)
				if err != io.EOF && !errors.Is(err, os.ErrClosed) && c.running.Load() {
					fmt.Printf("Error reading header: %v\n", err)
				}
				return
			}

			// Check for blank line (end of headers)
			if line == "\r\n" || line == "\n" {
				break
			}

			// Parse Content-Length
			var length int
			if _, err := fmt.Sscanf(line, "Content-Length: %d", &length); err == nil {
				contentLength = length
			}
		}

		if contentLength == 0 {
			continue
		}

		// Read message body
		body := make([]byte, contentLength)
		if _, err := io.ReadFull(reader, body); err != nil {
			// Only log unexpected errors (not EOF or closed pipe during shutdown)
			if err != io.EOF && !errors.Is(err, os.ErrClosed) && c.running.Load() {
				fmt.Printf("Error reading body: %v\n", err)
			}
			return
		}

		// Try to parse as request first (has both ID and Method)
		var request Request
		if err := json.Unmarshal(body, &request); err == nil && request.Method != "" {
			c.handleRequest(&request)
			continue
		}

		// Try to parse as response (has ID but no Method)
		var response Response
		if err := json.Unmarshal(body, &response); err == nil && len(response.ID) > 0 {
			c.handleResponse(&response)
			continue
		}
	}
}

// handleResponse dispatches a response to the waiting request
func (c *Client) handleResponse(response *Response) {
	var id string
	if err := json.Unmarshal(response.ID, &id); err != nil {
		return // ignore responses with non-string IDs
	}
	c.mu.Lock()
	responseChan, ok := c.pendingRequests[id]
	c.mu.Unlock()

	if ok {
		select {
		case responseChan <- response:
		default:
		}
	}
}

func (c *Client) handleRequest(request *Request) {
	c.mu.Lock()
	handler := c.requestHandlers[request.Method]
	c.mu.Unlock()

	if handler == nil {
		if request.IsCall() {
			c.sendErrorResponse(request.ID, -32601, fmt.Sprintf("Method not found: %s", request.Method), nil)
		}
		return
	}

	// Notifications run synchronously, calls run in a goroutine to avoid blocking
	if !request.IsCall() {
		handler(request.Params)
		return
	}

	go func() {
		defer func() {
			if r := recover(); r != nil {
				c.sendErrorResponse(request.ID, -32603, fmt.Sprintf("request handler panic: %v", r), nil)
			}
		}()

		result, err := handler(request.Params)
		if err != nil {
			c.sendErrorResponse(request.ID, err.Code, err.Message, err.Data)
			return
		}
		c.sendResponse(request.ID, result)
	}()
}

func (c *Client) sendResponse(id json.RawMessage, result json.RawMessage) {
	response := Response{
		JSONRPC: "2.0",
		ID:      id,
		Result:  result,
	}
	if err := c.sendMessage(response); err != nil {
		fmt.Printf("Failed to send JSON-RPC response: %v\n", err)
	}
}

func (c *Client) sendErrorResponse(id json.RawMessage, code int, message string, data map[string]any) {
	response := Response{
		JSONRPC: "2.0",
		ID:      id,
		Error: &Error{
			Code:    code,
			Message: message,
			Data:    data,
		},
	}
	if err := c.sendMessage(response); err != nil {
		fmt.Printf("Failed to send JSON-RPC error response: %v\n", err)
	}
}

// generateUUID generates a simple UUID v4 without external dependencies
func generateUUID() string {
	b := make([]byte, 16)
	rand.Read(b)
	b[6] = (b[6] & 0x0f) | 0x40 // Version 4
	b[8] = (b[8] & 0x3f) | 0x80 // Variant is 10
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}