package acp
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"io"
"strconv"
"sync"
"sync/atomic"
)
// Handler is the meaning behind the methods: one implementation per agent.
//
// Every method may be called concurrently, and Prompt in particular is
// expected to run for minutes while Cancel arrives on another goroutine. A
// handler that serializes its own state is a handler that works; this package
// deliberately does no locking on its behalf, because only the handler knows
// what a session is.
type Handler interface {
Initialize(context.Context, InitializeRequest) (InitializeResponse, error)
// Authenticate must not fail for an unknown method id. Clients send one
// unconditionally and do not catch the error, so a refusal here is a
// session that never starts rather than a login prompt.
Authenticate(context.Context, AuthenticateRequest) error
NewSession(context.Context, NewSessionRequest) (NewSessionResponse, error)
Prompt(context.Context, PromptRequest) (PromptResponse, error)
// Cancel is a notification: there is nothing to answer and nothing to
// fail. It must return promptly — the turn it cancels reports the outcome
// itself, by ending with StopCancelled.
Cancel(context.Context, CancelParams)
}
// Extender is the optional half of [Handler]: methods outside the core set.
// A handler that does not implement it answers CodeMethodNotFound to anything
// unrecognised, which is the correct answer and not an error.
type Extender interface {
Extra(ctx context.Context, method string, params json.RawMessage) (any, error)
}
// ErrClosed is what a pending [Conn.Request] fails with when the connection
// ends before the answer arrives.
var ErrClosed = errors.New("acp: connection closed")
// Conn is one ACP connection.
type Conn struct {
h Handler
r io.Reader
w io.Writer
// wmu serializes writes. Every frame is marshalled whole and written in
// one call, so interleaving is the only thing to guard against.
wmu sync.Mutex
nextID atomic.Int64
pmu sync.Mutex
pending map[string]chan *Message
closeOnce sync.Once
done chan struct{}
// wg counts handler goroutines, so Serve can wait for them to finish
// writing before it returns and the streams go away underneath them.
wg sync.WaitGroup
}
// NewConn wires a handler to a pair of streams. Nothing is read or written
// until [Conn.Serve] runs.
func NewConn(r io.Reader, w io.Writer, h Handler) *Conn {
return &Conn{
h: h,
r: r,
w: w,
pending: make(map[string]chan *Message),
done: make(chan struct{}),
}
}
// Serve reads until the stream ends, dispatching as it goes. A clean EOF is
// how a client says goodbye, so it returns nil; anything else is the error
// that stopped it.
//
// Cancelling ctx stops the handlers but not the read itself: a blocking read
// on an arbitrary io.Reader cannot be interrupted portably. Closing the reader
// is what ends Serve.
func (c *Conn) Serve(ctx context.Context) error {
ctx, cancel := context.WithCancel(ctx)
err := c.loop(ctx)
// Order matters. Cancel first so a turn in flight stops, then release
// anyone blocked waiting for an answer that is no longer coming, then wait
// for the handlers to finish with the streams.
cancel()
c.close()
c.wg.Wait()
return err
}
func (c *Conn) loop(ctx context.Context) error {
// ReadBytes rather than a Scanner: a prompt can carry a base64 image, and
// a line-length limit chosen here would be a decode failure nobody could
// explain.
br := bufio.NewReaderSize(c.r, 64*1024)
for {
line, err := br.ReadBytes('\n')
if len(bytes.TrimSpace(line)) > 0 {
c.handleLine(ctx, line)
}
switch {
case errors.Is(err, io.EOF):
return nil
case err != nil:
return err
}
}
}
func (c *Conn) handleLine(ctx context.Context, line []byte) {
var m Message
if err := json.Unmarshal(line, &m); err != nil {
// A frame we cannot parse might still have a recoverable id, in which
// case the peer is waiting and deserves an answer. If not, there is
// nobody to tell and the loop carries on: one bad line is not a reason
// to end a session.
if id := salvageID(line); id != nil {
c.respondError(id, Errorf(CodeParse, "invalid JSON"))
}
return
}
switch {
case m.ID != nil && m.Method == "":
c.deliver(&m)
case m.ID != nil:
c.wg.Add(1)
go func() {
defer c.wg.Done()
c.serveRequest(ctx, &m)
}()
case m.Method != "":
c.wg.Add(1)
go func() {
defer c.wg.Done()
c.serveNotification(ctx, &m)
}()
}
}
// salvageID pulls the id out of a line that did not parse as a whole, so a
// peer blocked on a request gets an error instead of a hang.
//
// It streams tokens rather than unmarshalling, because the interesting case is
// a frame that was cut off: the id is usually intact and everything after it
// is not, and a whole-value decode fails on exactly those lines.
func salvageID(line []byte) *json.RawMessage {
dec := json.NewDecoder(bytes.NewReader(line))
if tok, err := dec.Token(); err != nil || tok != json.Delim('{') {
return nil
}
for {
tok, err := dec.Token()
if err != nil {
return nil
}
key, ok := tok.(string)
if !ok {
return nil // the closing brace, or something that is not an object
}
var raw json.RawMessage
if err := dec.Decode(&raw); err != nil {
return nil
}
if key == "id" {
return &raw
}
}
}
func (c *Conn) serveRequest(ctx context.Context, m *Message) {
result, err := c.call(ctx, m.Method, m.Params)
if err != nil {
c.respondError(m.ID, asError(err))
return
}
c.respond(m.ID, result)
}
func (c *Conn) serveNotification(ctx context.Context, m *Message) {
if m.Method != MethodCancel {
// Unknown notifications are ignored by design: that is what makes it
// safe for a client to send one we have not heard of.
return
}
var params CancelParams
if err := decodeParams(m.Params, ¶ms); err != nil {
return
}
c.h.Cancel(ctx, params)
}
// call dispatches one request to the handler. The returned value is whatever
// goes in the result field; nil becomes an empty object, because a JSON-RPC
// response with no result at all is not a success.
func (c *Conn) call(ctx context.Context, method string, params json.RawMessage) (any, error) {
switch method {
case MethodInitialize:
var req InitializeRequest
if err := decodeParams(params, &req); err != nil {
return nil, err
}
return c.h.Initialize(ctx, req)
case MethodAuthenticate:
var req AuthenticateRequest
if err := decodeParams(params, &req); err != nil {
return nil, err
}
return nil, c.h.Authenticate(ctx, req)
case MethodNewSession:
var req NewSessionRequest
if err := decodeParams(params, &req); err != nil {
return nil, err
}
return c.h.NewSession(ctx, req)
case MethodPrompt:
var req PromptRequest
if err := decodeParams(params, &req); err != nil {
return nil, err
}
return c.h.Prompt(ctx, req)
default:
if ex, ok := c.h.(Extender); ok {
return ex.Extra(ctx, method, params)
}
return nil, Errorf(CodeMethodNotFound, "unknown method %q", method)
}
}
// Notify sends a notification. It never has an id, which is what stops the
// peer from waiting for an answer to it.
func (c *Conn) Notify(ctx context.Context, method string, params any) error {
raw, err := json.Marshal(params)
if err != nil {
return err
}
return c.write(&Message{JSONRPC: jsonrpcVersion, Method: method, Params: raw})
}
// Request calls the peer and waits for the answer, decoding it into out when
// out is non-nil. It returns ctx.Err() if the caller gives up first and
// [ErrClosed] if the connection ends before the peer replies — a permission
// request outliving its connection has to fail rather than hang the turn.
func (c *Conn) Request(ctx context.Context, method string, params, out any) error {
raw, err := json.Marshal(params)
if err != nil {
return err
}
id := json.RawMessage(strconv.FormatInt(c.nextID.Add(1), 10))
key := string(id)
ch := make(chan *Message, 1)
c.pmu.Lock()
c.pending[key] = ch
c.pmu.Unlock()
defer func() {
c.pmu.Lock()
delete(c.pending, key)
c.pmu.Unlock()
}()
if err := c.write(&Message{JSONRPC: jsonrpcVersion, ID: &id, Method: method, Params: raw}); err != nil {
return err
}
select {
case <-ctx.Done():
return ctx.Err()
case <-c.done:
return ErrClosed
case reply := <-ch:
if reply.Error != nil {
return reply.Error
}
if out == nil || len(reply.Result) == 0 {
return nil
}
return json.Unmarshal(reply.Result, out)
}
}
func (c *Conn) deliver(m *Message) {
c.pmu.Lock()
ch := c.pending[string(*m.ID)]
c.pmu.Unlock()
if ch != nil {
// Buffered, and the entry is removed by the waiter, so this cannot
// block and cannot be delivered twice.
select {
case ch <- m:
default:
}
}
}
func (c *Conn) respond(id *json.RawMessage, result any) {
if result == nil {
result = struct{}{}
}
raw, err := json.Marshal(result)
if err != nil {
c.respondError(id, asError(err))
return
}
_ = c.write(&Message{JSONRPC: jsonrpcVersion, ID: id, Result: raw})
}
func (c *Conn) respondError(id *json.RawMessage, e *Error) {
if id == nil {
null := json.RawMessage("null")
id = &null
}
_ = c.write(&Message{JSONRPC: jsonrpcVersion, ID: id, Error: e})
}
func (c *Conn) write(m *Message) error {
b, err := json.Marshal(m)
if err != nil {
return err
}
b = append(b, '\n')
c.wmu.Lock()
defer c.wmu.Unlock()
_, err = c.w.Write(b)
return err
}
func (c *Conn) close() {
c.closeOnce.Do(func() { close(c.done) })
}