1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101

package acp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
)

// MaxMessageBytes caps how large a single message may be. An agent should
// never send anything near this; the limit is there so that a stream with no
// newline in it cannot make the editor allocate without bound.
const MaxMessageBytes = 64 << 20

// ErrMessageTooLarge is returned when a line runs past MaxMessageBytes.
var ErrMessageTooLarge = errors.New("acp: message too large")

// Framing is the Agent Client Protocol's way of marking one message off from
// the next: a newline. The specification puts it plainly — messages are
// delimited by "\n" and must not contain embedded newlines.
//
// It is the jsonrpc.Framer an agent connection is built with.
//
//	conn := jsonrpc.NewConn(stream, acp.Framing{}, onNotify, onRequest)
type Framing struct{}

// ReadMessage reads one line and returns it, without the newline.
//
// It returns io.EOF when the stream ends cleanly between messages, which is how
// an agent shutting down is told apart from one that died mid-message.
//
// Blank lines are skipped rather than handed on as empty messages: an agent
// that prints a stray newline should not end the conversation.
func (Framing) ReadMessage(r *bufio.Reader) ([]byte, error) {
	for {
		line, err := readLine(r)
		if err != nil {
			return nil, err
		}
		if len(line) > 0 {
			return line, nil
		}
	}
}

// readLine reads up to one newline, refusing a line longer than the cap.
//
// bufio's own ReadBytes would grow without bound, which is exactly what the cap
// exists to prevent — a peer writing megabytes with no newline in them.
func readLine(r *bufio.Reader) ([]byte, error) {
	var line []byte
	for {
		chunk, err := r.ReadSlice('\n')
		line = append(line, chunk...)

		switch {
		case err == nil:
			return trimNewline(line), nil
		case errors.Is(err, bufio.ErrBufferFull):
			if len(line) > MaxMessageBytes {
				return nil, fmt.Errorf("%w: over %d bytes with no newline", ErrMessageTooLarge, MaxMessageBytes)
			}
		case errors.Is(err, io.EOF):
			// A last line with no newline after it is still a message; an
			// agent killed mid-write leaves one, and so does a fixture file.
			if trimmed := trimNewline(line); len(trimmed) > 0 {
				return trimmed, nil
			}
			return nil, io.EOF
		default:
			return nil, err
		}
	}
}

// trimNewline drops the line ending, tolerating the \r\n an agent on Windows
// may write.
func trimNewline(line []byte) []byte {
	line = dropSuffix(line, '\n')
	return dropSuffix(line, '\r')
}

// dropSuffix removes one trailing byte if it is there.
func dropSuffix(line []byte, b byte) []byte {
	if len(line) > 0 && line[len(line)-1] == b {
		return line[:len(line)-1]
	}
	return line
}

// WriteMessage writes body followed by a newline.
//
// Nothing escapes body first: the only caller is jsonrpc, which passes
// json.Marshal's output, and that never contains a literal newline — the
// encoder writes \n inside a string as the two characters backslash and n.
func (Framing) WriteMessage(w io.Writer, body []byte) error {
	if _, err := w.Write(append(body, '\n')); err != nil {
		return fmt.Errorf("acp: writing a message: %w", err)
	}
	return nil
}