| 🛟 Updated. 28d5985 k33g 20h ago | 1 | package acp |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "errors" |
| 6 | "fmt" |
| 7 | "io" |
| 8 | ) |
| 9 | |
| 10 | // MaxMessageBytes caps how large a single message may be. An agent should |
| 11 | // never send anything near this; the limit is there so that a stream with no |
| 12 | // newline in it cannot make the editor allocate without bound. |
| 13 | const MaxMessageBytes = 64 << 20 |
| 14 | |
| 15 | // ErrMessageTooLarge is returned when a line runs past MaxMessageBytes. |
| 16 | var ErrMessageTooLarge = errors.New("acp: message too large") |
| 17 | |
| 18 | // Framing is the Agent Client Protocol's way of marking one message off from |
| 19 | // the next: a newline. The specification puts it plainly — messages are |
| 20 | // delimited by "\n" and must not contain embedded newlines. |
| 21 | // |
| 22 | // It is the jsonrpc.Framer an agent connection is built with. |
| 23 | // |
| 24 | // conn := jsonrpc.NewConn(stream, acp.Framing{}, onNotify, onRequest) |
| 25 | type Framing struct{} |
| 26 | |
| 27 | // ReadMessage reads one line and returns it, without the newline. |
| 28 | // |
| 29 | // It returns io.EOF when the stream ends cleanly between messages, which is how |
| 30 | // an agent shutting down is told apart from one that died mid-message. |
| 31 | // |
| 32 | // Blank lines are skipped rather than handed on as empty messages: an agent |
| 33 | // that prints a stray newline should not end the conversation. |
| 34 | func (Framing) ReadMessage(r *bufio.Reader) ([]byte, error) { |
| 35 | for { |
| 36 | line, err := readLine(r) |
| 37 | if err != nil { |
| 38 | return nil, err |
| 39 | } |
| 40 | if len(line) > 0 { |
| 41 | return line, nil |
| 42 | } |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | // readLine reads up to one newline, refusing a line longer than the cap. |
| 47 | // |
| 48 | // bufio's own ReadBytes would grow without bound, which is exactly what the cap |
| 49 | // exists to prevent — a peer writing megabytes with no newline in them. |
| 50 | func readLine(r *bufio.Reader) ([]byte, error) { |
| 51 | var line []byte |
| 52 | for { |
| 53 | chunk, err := r.ReadSlice('\n') |
| 54 | line = append(line, chunk...) |
| 55 | |
| 56 | switch { |
| 57 | case err == nil: |
| 58 | return trimNewline(line), nil |
| 59 | case errors.Is(err, bufio.ErrBufferFull): |
| 60 | if len(line) > MaxMessageBytes { |
| 61 | return nil, fmt.Errorf("%w: over %d bytes with no newline", ErrMessageTooLarge, MaxMessageBytes) |
| 62 | } |
| 63 | case errors.Is(err, io.EOF): |
| 64 | // A last line with no newline after it is still a message; an |
| 65 | // agent killed mid-write leaves one, and so does a fixture file. |
| 66 | if trimmed := trimNewline(line); len(trimmed) > 0 { |
| 67 | return trimmed, nil |
| 68 | } |
| 69 | return nil, io.EOF |
| 70 | default: |
| 71 | return nil, err |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | // trimNewline drops the line ending, tolerating the \r\n an agent on Windows |
| 77 | // may write. |
| 78 | func trimNewline(line []byte) []byte { |
| 79 | line = dropSuffix(line, '\n') |
| 80 | return dropSuffix(line, '\r') |
| 81 | } |
| 82 | |
| 83 | // dropSuffix removes one trailing byte if it is there. |
| 84 | func dropSuffix(line []byte, b byte) []byte { |
| 85 | if len(line) > 0 && line[len(line)-1] == b { |
| 86 | return line[:len(line)-1] |
| 87 | } |
| 88 | return line |
| 89 | } |
| 90 | |
| 91 | // WriteMessage writes body followed by a newline. |
| 92 | // |
| 93 | // Nothing escapes body first: the only caller is jsonrpc, which passes |
| 94 | // json.Marshal's output, and that never contains a literal newline — the |
| 95 | // encoder writes \n inside a string as the two characters backslash and n. |
| 96 | func (Framing) WriteMessage(w io.Writer, body []byte) error { |
| 97 | if _, err := w.Write(append(body, '\n')); err != nil { |
| 98 | return fmt.Errorf("acp: writing a message: %w", err) |
| 99 | } |
| 100 | return nil |
| 101 | } |