codespace/internal/acp/transport.go

package acp

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"sync"
	"sync/atomic"

	"codespace/internal/process"
)

// transport reads NDJSON frames from a process subscription and writes NDJSON
// frames to the process stdin.
type transport struct {
	stdin          io.WriteCloser
	sub            process.Subscription
	frames         chan *jsonRPCMessage
	pending        map[int]chan *jsonRPCMessage
	pendingMu      sync.Mutex
	nextID         atomic.Int64
	onRequest      func(method string, params json.RawMessage, id int)
	onNotification func(method string, params json.RawMessage)
	onClose        func()
	lg             *slog.Logger
	closeOnce      sync.Once
}

// newTransport creates a transport bound to the given stdin writer and output
// subscription. Call start to begin the read/dispatch goroutines.
func newTransport(stdin io.WriteCloser, sub process.Subscription, lg *slog.Logger) *transport {
	if lg == nil {
		lg = slog.Default()
	}
	return &transport{
		stdin:   stdin,
		sub:     sub,
		frames:  make(chan *jsonRPCMessage, 64),
		pending: make(map[int]chan *jsonRPCMessage),
		lg:      lg,
	}
}

// start launches the read and dispatch loops.
func (t *transport) start() {
	go t.readLoop()
	go t.dispatchLoop()
}

// readLoop consumes raw bytes from the subscription, splits on newlines, and
// feeds parsed JSON-RPC frames to the frames channel.
func (t *transport) readLoop() {
	defer close(t.frames)
	buf := bytes.NewBuffer(make([]byte, 0, 4096))
	for chunk := range t.sub.Output() {
		buf.Write(chunk)
		for {
			data := buf.Bytes()
			idx := bytes.IndexByte(data, '\n')
			if idx < 0 {
				break
			}
			line := make([]byte, idx)
			copy(line, data[:idx])
			buf.Next(idx + 1)
			msg := &jsonRPCMessage{}
			if err := json.Unmarshal(line, msg); err != nil {
				t.lg.Error("acp invalid ndjson line", "error", err)
				continue
			}
			t.frames <- msg
		}
	}
}

// dispatchLoop routes frames to pending request channels, notification
// handlers, or incoming-request handlers.
func (t *transport) dispatchLoop() {
	defer func() {
		if t.onClose != nil {
			t.onClose()
		}
	}()
	for msg := range t.frames {
		switch {
		case msg.ID != nil && msg.Method != "":
			// Incoming request from the agent.
			if t.onRequest != nil {
				id := *msg.ID
				go t.onRequest(msg.Method, msg.Params, id)
			}
		case msg.ID == nil && msg.Method != "":
			// Notification from the agent.
			if t.onNotification != nil {
				go t.onNotification(msg.Method, msg.Params)
			}
		case msg.ID != nil && (msg.Result != nil || msg.Error != nil):
			// Response to a client-initiated request.
			id := *msg.ID
			t.pendingMu.Lock()
			ch, ok := t.pending[id]
			delete(t.pending, id)
			t.pendingMu.Unlock()
			if ok {
				select {
				case ch <- msg:
				default:
				}
			} else {
				t.lg.Debug("acp unsolicited response", "id", id)
			}
		default:
			t.lg.Debug("acp unhandled frame", "method", msg.Method, "id", msg.ID)
		}
	}
}

// call sends a JSON-RPC request and waits for its response.
func (t *transport) call(ctx context.Context, method string, params, result any) error {
	id := int(t.nextID.Add(1))
	rawParams, err := json.Marshal(params)
	if err != nil {
		return err
	}
	msg := &jsonRPCMessage{JSONRPC: "2.0", ID: &id, Method: method, Params: rawParams}
	ch := make(chan *jsonRPCMessage, 1)
	t.pendingMu.Lock()
	t.pending[id] = ch
	t.pendingMu.Unlock()
	defer func() {
		t.pendingMu.Lock()
		delete(t.pending, id)
		t.pendingMu.Unlock()
	}()
	if err := t.write(msg); err != nil {
		return err
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case resp := <-ch:
		if resp.Error != nil {
			return errors.New(resp.Error.Message)
		}
		if result != nil && resp.Result != nil {
			return json.Unmarshal(resp.Result, result)
		}
		return nil
	}
}

// sendNotification writes a JSON-RPC notification (no id).
func (t *transport) sendNotification(method string, params any) error {
	rawParams, err := json.Marshal(params)
	if err != nil {
		return err
	}
	msg := &jsonRPCMessage{JSONRPC: "2.0", Method: method, Params: rawParams}
	return t.write(msg)
}

// sendResponse writes a JSON-RPC response for an agent-initiated request.
func (t *transport) sendResponse(id int, result any, rpcErr *rpcError) error {
	msg := &jsonRPCMessage{JSONRPC: "2.0", ID: &id}
	if rpcErr != nil {
		msg.Error = rpcErr
	} else {
		raw, err := json.Marshal(result)
		if err != nil {
			return err
		}
		msg.Result = raw
	}
	return t.write(msg)
}

// write marshals a JSON-RPC message and appends a newline.
func (t *transport) write(msg *jsonRPCMessage) error {
	data, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	data = append(data, '\n')
	_, err = t.stdin.Write(data)
	return err
}

// Close shuts down the transport by closing the subscription. This causes the
// read loop to exit and the dispatch loop to follow.
func (t *transport) Close() error {
	t.closeOnce.Do(func() {
		_ = t.sub.Close()
	})
	return nil
}