| 1 | package acp |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "errors" |
| 8 | "io" |
| 9 | "strconv" |
| 10 | "sync" |
| 11 | "sync/atomic" |
| 12 | ) |
| 13 | |
| 14 | // maxMessageBytes caps a single inbound NDJSON line. ACP messages can embed |
| 15 | // resource text, so the limit is generous; a line past it is a framing error. |
| 16 | const maxMessageBytes = 32 << 20 // 32 MiB |
| 17 | |
| 18 | // RequestHandler answers an inbound JSON-RPC request. The returned value is |
| 19 | // marshaled as the response result. To control the error code, return a |
| 20 | // *RPCError; any other error becomes ErrInternal. |
| 21 | type RequestHandler func(ctx context.Context, params json.RawMessage) (any, error) |
| 22 | |
| 23 | // responseWithAfter lets a request handler schedule work that must run only |
| 24 | // after its successful JSON-RPC response has been written to the wire. |
| 25 | type responseWithAfter interface { |
| 26 | Response() any |
| 27 | AfterResponse() |
| 28 | } |
| 29 | |
| 30 | // NotificationHandler reacts to an inbound notification. It cannot reply, so it |
| 31 | // returns nothing — errors have nowhere to go on the wire (stderr would corrupt |
| 32 | // stdout, which is the JSON-RPC channel). |
| 33 | type NotificationHandler func(ctx context.Context, params json.RawMessage) |
| 34 | |
| 35 | // RPCError lets a handler choose the JSON-RPC error code returned to the client. |
| 36 | type RPCError struct { |
| 37 | Code int |
| 38 | Message string |
| 39 | } |
| 40 | |
| 41 | func (e *RPCError) Error() string { return e.Message } |
| 42 | |
| 43 | // Conn is one NDJSON JSON-RPC 2.0 connection over a reader/writer pair (stdin/ |
| 44 | // stdout in production). It dispatches inbound requests and notifications to |
| 45 | // registered handlers, and can itself send outbound notifications (session/update) |
| 46 | // and requests (session/request_permission), correlating replies by id. |
| 47 | // |
| 48 | // Writes are serialized by a mutex, so handlers running on separate goroutines |
| 49 | // (a long session/prompt alongside a session/cancel) never interleave a line. |
| 50 | // It implements notifier, the dependency the dispatch sink takes. |
| 51 | type Conn struct { |
| 52 | r io.Reader |
| 53 | |
| 54 | wmu sync.Mutex |
| 55 | enc *json.Encoder |
| 56 | |
| 57 | nextID atomic.Int64 |
| 58 | |
| 59 | pmu sync.Mutex |
| 60 | pending map[int64]chan rpcResult |
| 61 | |
| 62 | reqH map[string]RequestHandler |
| 63 | notH map[string]NotificationHandler |
| 64 | |
| 65 | wg sync.WaitGroup |
| 66 | closeOnce sync.Once |
| 67 | closed chan struct{} |
| 68 | } |
| 69 | |
| 70 | // rpcResult is the outcome of an outbound request, delivered to the waiter. |
| 71 | type rpcResult struct { |
| 72 | result json.RawMessage |
| 73 | err error |
| 74 | } |
| 75 | |
| 76 | // rpcError is the JSON-RPC error object on the wire. |
| 77 | type rpcError struct { |
| 78 | Code int `json:"code"` |
| 79 | Message string `json:"message"` |
| 80 | Data json.RawMessage `json:"data,omitempty"` |
| 81 | } |
| 82 | |
| 83 | // outbound is a JSON-RPC frame we send. omitempty fields select between request, |
| 84 | // notification, success response, and error response shapes. |
| 85 | type outbound struct { |
| 86 | JSONRPC string `json:"jsonrpc"` |
| 87 | ID json.RawMessage `json:"id,omitempty"` |
| 88 | Method string `json:"method,omitempty"` |
| 89 | Params json.RawMessage `json:"params,omitempty"` |
| 90 | Result json.RawMessage `json:"result,omitempty"` |
| 91 | Error *rpcError `json:"error,omitempty"` |
| 92 | } |
| 93 | |
| 94 | // inbound is a parsed JSON-RPC frame we received. The combination of id/method |
| 95 | // presence distinguishes request, notification, and response. |
| 96 | type inbound struct { |
| 97 | ID json.RawMessage `json:"id"` |
| 98 | Method string `json:"method"` |
| 99 | Params json.RawMessage `json:"params"` |
| 100 | Result json.RawMessage `json:"result"` |
| 101 | Error *rpcError `json:"error"` |
| 102 | } |
| 103 | |
| 104 | // NewConn wires a connection over r (inbound) and w (outbound). Register handlers |
| 105 | // with Handle / HandleNotify before calling Serve. The encoder disables HTML |
| 106 | // escaping so payloads match main's JSON.stringify output byte-for-byte. |
| 107 | func NewConn(r io.Reader, w io.Writer) *Conn { |
| 108 | enc := json.NewEncoder(w) |
| 109 | enc.SetEscapeHTML(false) |
| 110 | return &Conn{ |
| 111 | r: r, |
| 112 | enc: enc, |
| 113 | pending: make(map[int64]chan rpcResult), |
| 114 | reqH: make(map[string]RequestHandler), |
| 115 | notH: make(map[string]NotificationHandler), |
| 116 | closed: make(chan struct{}), |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | // Handle registers a request handler for method. Not safe to call concurrently |
| 121 | // with Serve; wire all handlers up first. |
| 122 | func (c *Conn) Handle(method string, h RequestHandler) { c.reqH[method] = h } |
| 123 | |
| 124 | // HandleNotify registers a notification handler for method. |
| 125 | func (c *Conn) HandleNotify(method string, h NotificationHandler) { c.notH[method] = h } |
| 126 | |
| 127 | // Serve reads inbound frames until the reader ends or ctx is cancelled. Each |
| 128 | // inbound request/notification runs on its own goroutine so a long-running prompt |
| 129 | // does not block cancellation or permission replies. When the read loop ends it |
| 130 | // cancels in-flight handlers (so prompts abort) and waits for them to return — |
| 131 | // flushing fast responses and unwinding aborted ones — before failing any |
| 132 | // outstanding outbound requests. Returns nil on clean EOF. |
| 133 | func (c *Conn) Serve(ctx context.Context) error { |
| 134 | ctx, cancel := context.WithCancel(ctx) |
| 135 | defer cancel() |
| 136 | |
| 137 | br := bufio.NewReaderSize(c.r, 64<<10) |
| 138 | var loopErr error |
| 139 | for { |
| 140 | line, err := readLine(br) |
| 141 | if len(line) > 0 { |
| 142 | c.dispatch(ctx, line) |
| 143 | } |
| 144 | if err != nil { |
| 145 | if !errors.Is(err, io.EOF) { |
| 146 | loopErr = err |
| 147 | } |
| 148 | break |
| 149 | } |
| 150 | if err := ctx.Err(); err != nil { |
| 151 | loopErr = err |
| 152 | break |
| 153 | } |
| 154 | } |
| 155 | |
| 156 | cancel() // abort in-flight handlers (prompts unwind via ctx) |
| 157 | c.wg.Wait() // let them flush their responses before we tear down |
| 158 | c.shutdown() // fail any still-pending outbound requests |
| 159 | return loopErr |
| 160 | } |
| 161 | |
| 162 | // readLine reads one NDJSON line (without the trailing newline), enforcing the |
| 163 | // size cap. It returns the line and any read error; on EOF it still returns the |
| 164 | // trailing partial line so a final newline-less frame is processed. |
| 165 | func readLine(br *bufio.Reader) ([]byte, error) { |
| 166 | var buf []byte |
| 167 | for { |
| 168 | chunk, err := br.ReadSlice('\n') |
| 169 | buf = append(buf, chunk...) |
| 170 | if len(buf) > maxMessageBytes { |
| 171 | return nil, errors.New("acp: message exceeds size limit") |
| 172 | } |
| 173 | if errors.Is(err, bufio.ErrBufferFull) { |
| 174 | continue |
| 175 | } |
| 176 | // Trim the trailing newline (and CR) if present. |
| 177 | n := len(buf) |
| 178 | for n > 0 && (buf[n-1] == '\n' || buf[n-1] == '\r') { |
| 179 | n-- |
| 180 | } |
| 181 | return trimSpace(buf[:n]), err |
| 182 | } |
| 183 | } |
| 184 | |
| 185 | // trimSpace drops leading/trailing ASCII whitespace without allocating. |
| 186 | func trimSpace(b []byte) []byte { |
| 187 | i, j := 0, len(b) |
| 188 | for i < j && isSpace(b[i]) { |
| 189 | i++ |
| 190 | } |
| 191 | for j > i && isSpace(b[j-1]) { |
| 192 | j-- |
| 193 | } |
| 194 | return b[i:j] |
| 195 | } |
| 196 | |
| 197 | func isSpace(c byte) bool { return c == ' ' || c == '\t' || c == '\n' || c == '\r' } |
| 198 | |
| 199 | // dispatch parses one frame and routes it. Requests and notifications fan out to |
| 200 | // goroutines; responses resolve inline (they are cheap and need ordering only |
| 201 | // against the pending map, which is mutex-guarded). |
| 202 | func (c *Conn) dispatch(ctx context.Context, line []byte) { |
| 203 | var in inbound |
| 204 | if err := json.Unmarshal(line, &in); err != nil { |
| 205 | c.writeError(json.RawMessage("null"), ErrParse, "parse error") |
| 206 | return |
| 207 | } |
| 208 | hasID := len(in.ID) > 0 |
| 209 | switch { |
| 210 | case in.Method != "" && hasID: |
| 211 | c.wg.Go(func() { |
| 212 | c.serveRequest(ctx, in.ID, in.Method, in.Params) |
| 213 | }) |
| 214 | case in.Method != "" && !hasID: |
| 215 | if h := c.notH[in.Method]; h != nil { |
| 216 | c.wg.Go(func() { |
| 217 | h(ctx, in.Params) |
| 218 | }) |
| 219 | } |
| 220 | case in.Method == "" && hasID: |
| 221 | c.resolve(in) |
| 222 | default: |
| 223 | c.writeError(json.RawMessage("null"), ErrInvalidRequest, "invalid request") |
| 224 | } |
| 225 | } |
| 226 | |
| 227 | // serveRequest runs a request handler and writes its response (or error). |
| 228 | func (c *Conn) serveRequest(ctx context.Context, id json.RawMessage, method string, params json.RawMessage) { |
| 229 | h := c.reqH[method] |
| 230 | if h == nil { |
| 231 | c.writeError(id, ErrMethodNotFound, "method not found: "+method) |
| 232 | return |
| 233 | } |
| 234 | result, err := h(ctx, params) |
| 235 | if err != nil { |
| 236 | code := ErrInternal |
| 237 | var re *RPCError |
| 238 | if errors.As(err, &re) { |
| 239 | code = re.Code |
| 240 | } |
| 241 | c.writeError(id, code, err.Error()) |
| 242 | return |
| 243 | } |
| 244 | var after func() |
| 245 | if r, ok := result.(responseWithAfter); ok { |
| 246 | result = r.Response() |
| 247 | after = r.AfterResponse |
| 248 | } |
| 249 | raw, err := json.Marshal(result) |
| 250 | if err != nil { |
| 251 | c.writeError(id, ErrInternal, "marshal result: "+err.Error()) |
| 252 | return |
| 253 | } |
| 254 | if err := c.write(outbound{JSONRPC: "2.0", ID: id, Result: raw}); err != nil { |
| 255 | return |
| 256 | } |
| 257 | if after != nil { |
| 258 | after() |
| 259 | } |
| 260 | } |
| 261 | |
| 262 | // resolve delivers a response to the goroutine waiting on its outbound request. |
| 263 | func (c *Conn) resolve(in inbound) { |
| 264 | id, err := strconv.ParseInt(string(in.ID), 10, 64) |
| 265 | if err != nil { |
| 266 | return // we only issue integer ids; an unparsable id isn't ours |
| 267 | } |
| 268 | c.pmu.Lock() |
| 269 | ch := c.pending[id] |
| 270 | delete(c.pending, id) |
| 271 | c.pmu.Unlock() |
| 272 | if ch == nil { |
| 273 | return |
| 274 | } |
| 275 | if in.Error != nil { |
| 276 | ch <- rpcResult{err: errors.New(in.Error.Message)} |
| 277 | return |
| 278 | } |
| 279 | ch <- rpcResult{result: in.Result} |
| 280 | } |
| 281 | |
| 282 | // Notify sends a fire-and-forget notification. Satisfies notifier. |
| 283 | func (c *Conn) Notify(method string, params any) error { |
| 284 | raw, err := json.Marshal(params) |
| 285 | if err != nil { |
| 286 | return err |
| 287 | } |
| 288 | return c.write(outbound{JSONRPC: "2.0", Method: method, Params: raw}) |
| 289 | } |
| 290 | |
| 291 | // Request sends an outbound request and blocks until the peer responds, ctx is |
| 292 | // cancelled, or the connection closes. Satisfies notifier. |
| 293 | func (c *Conn) Request(ctx context.Context, method string, params any) (json.RawMessage, error) { |
| 294 | raw, err := json.Marshal(params) |
| 295 | if err != nil { |
| 296 | return nil, err |
| 297 | } |
| 298 | id := c.nextID.Add(1) |
| 299 | ch := make(chan rpcResult, 1) |
| 300 | c.pmu.Lock() |
| 301 | c.pending[id] = ch |
| 302 | c.pmu.Unlock() |
| 303 | defer func() { |
| 304 | c.pmu.Lock() |
| 305 | delete(c.pending, id) |
| 306 | c.pmu.Unlock() |
| 307 | }() |
| 308 | |
| 309 | idRaw, _ := json.Marshal(id) |
| 310 | if err := c.write(outbound{JSONRPC: "2.0", ID: idRaw, Method: method, Params: raw}); err != nil { |
| 311 | return nil, err |
| 312 | } |
| 313 | select { |
| 314 | case res := <-ch: |
| 315 | return res.result, res.err |
| 316 | case <-ctx.Done(): |
| 317 | return nil, ctx.Err() |
| 318 | case <-c.closed: |
| 319 | return nil, errors.New("acp: connection closed") |
| 320 | } |
| 321 | } |
| 322 | |
| 323 | func (c *Conn) write(m outbound) error { |
| 324 | c.wmu.Lock() |
| 325 | defer c.wmu.Unlock() |
| 326 | return c.enc.Encode(m) |
| 327 | } |
| 328 | |
| 329 | func (c *Conn) writeError(id json.RawMessage, code int, msg string) { |
| 330 | _ = c.write(outbound{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: code, Message: msg}}) |
| 331 | } |
| 332 | |
| 333 | // shutdown fails every in-flight outbound request so their goroutines unblock. |
| 334 | func (c *Conn) shutdown() { |
| 335 | c.closeOnce.Do(func() { |
| 336 | close(c.closed) |
| 337 | c.pmu.Lock() |
| 338 | for id, ch := range c.pending { |
| 339 | ch <- rpcResult{err: errors.New("acp: connection closed")} |
| 340 | delete(c.pending, id) |
| 341 | } |
| 342 | c.pmu.Unlock() |
| 343 | }) |
| 344 | } |
| 345 |