From 3d5f63b28e64d021768f2e25d402ee8dedbdc675 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Thu, 7 May 2026 12:25:39 +0800 Subject: [PATCH] feat(acp): relay reader/writer goroutines + Call + persist with truncation --- internal/acp/relay.go | 225 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 225 insertions(+) diff --git a/internal/acp/relay.go b/internal/acp/relay.go index e288c2c..bd9c155 100644 --- a/internal/acp/relay.go +++ b/internal/acp/relay.go @@ -10,8 +10,13 @@ package acp import ( + "context" + "crypto/sha256" + "encoding/hex" "encoding/json" + "errors" "fmt" + "io" "log/slog" "sync" "sync/atomic" @@ -187,3 +192,223 @@ func mustJSON(v any) json.RawMessage { } return b } + +// Run is the relay main loop; supervisor.Spawn calls it in a goroutine. +// It launches the reader and writer goroutines and waits for both. +func (r *Relay) Run(ctx context.Context, sess handlers.SessionContext) { + var wg sync.WaitGroup + wg.Add(2) + go func() { defer wg.Done(); r.reader(ctx, sess) }() + go func() { defer wg.Done(); r.writer(ctx) }() + wg.Wait() +} + +func (r *Relay) reader(ctx context.Context, sess handlers.SessionContext) { + for { + msg, err := r.dec.DecodeMessage() + if err != nil { + if errors.Is(err, io.EOF) { + r.log.Info("acp.relay.eof", "session_id", r.sessionID) + return + } + r.log.Error("acp.relay.decode", "session_id", r.sessionID, "err", err.Error()) + return + } + + envelope := r.persistAndEnvelope(ctx, msg, "out") + if envelope != nil { + r.fanout(envelope) + } + + switch { + case msg.IsResponse(): + if id, ok := msg.IDInt64(); ok { + if ch, hit := r.inflight.LoadAndDelete(id); hit { + ch.(chan *Message) <- msg + } else { + r.log.Warn("acp.relay.orphan_response", "session_id", r.sessionID, "id", id) + } + } + // String IDs go to clients via fanout (handled above); no server action. + case msg.IsNotification(): + // Already persisted + fanout'd. No reply. + case msg.IsRequest(): + go r.handleAgentRequest(ctx, sess, msg) + } + } +} + +// handleAgentRequest dispatches an agent → server request to the appropriate handler. +func (r *Relay) handleAgentRequest(ctx context.Context, sess handlers.SessionContext, req *Message) { + var resp *Message + switch req.Method { + case "fs/read_text_file": + out := r.fsHandler.Read.Handle(ctx, sess, req.Params) + resp = r.fsResultToResponse(req.ID, out) + case "fs/write_text_file": + out := r.fsHandler.Write.Handle(ctx, sess, req.Params) + resp = r.fsResultToResponse(req.ID, out) + case "session/request_permission": + out := r.permHandler.Handle(ctx, sess, req.Params) + resp = r.fsResultToResponse(req.ID, out) + default: + resp = NewResponseErr(req.ID, JSONRPCMethodNotFound, "method not implemented: "+req.Method) + } + select { + case r.serverInitiated <- resp: + case <-ctx.Done(): + } +} + +func (r *Relay) fsResultToResponse(id json.RawMessage, res handlers.FsResult) *Message { + if res.Err != nil { + return NewResponseErr(id, res.Err.Code, res.Err.Message) + } + return &Message{JSONRPC: "2.0", ID: id, Result: res.OK} +} + +func (r *Relay) writer(ctx context.Context) { + for { + select { + case <-ctx.Done(): + return + case <-r.done: + return + case msg := <-r.serverInitiated: + r.persistAndEnvelope(ctx, msg, "in") + if err := r.enc.Encode(msg); err != nil { + r.log.Error("acp.relay.encode", "session_id", r.sessionID, "err", err.Error()) + return + } + case msg := <-r.clientInitiated: + r.persistAndEnvelope(ctx, msg, "in") + if err := r.enc.Encode(msg); err != nil { + r.log.Error("acp.relay.encode", "session_id", r.sessionID, "err", err.Error()) + return + } + } + } +} + +// SendClient is the WS handler entry point for forwarding prompt/cancel from a tab. +func (r *Relay) SendClient(msg *Message) error { + if r.closed.Load() { + return fmt.Errorf("relay closed") + } + select { + case r.clientInitiated <- msg: + return nil + default: + return fmt.Errorf("relay client channel full") + } +} + +// Call performs a server-initiated request and waits for the response (e.g. +// initialize / session/new). Used by supervisor during handshake. +func (r *Relay) Call(ctx context.Context, method string, params any) (*Message, error) { + id := r.nextID.Add(1) + req, err := NewRequestInt(id, method, params) + if err != nil { + return nil, err + } + ch := make(chan *Message, 1) + r.inflight.Store(id, ch) + defer r.inflight.Delete(id) + + select { + case r.serverInitiated <- req: + case <-ctx.Done(): + return nil, ctx.Err() + } + + select { + case resp := <-ch: + if resp.Error != nil { + return resp, resp.Error + } + return resp, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// persistAndEnvelope writes one event to the repo and returns the envelope for +// fanout. Returns nil if the write fails (logged). +func (r *Relay) persistAndEnvelope(ctx context.Context, msg *Message, direction string) *EventEnvelope { + rpcKind := classifyKind(msg) + method := msg.Method + payload, _ := json.Marshal(msg) + + origSize := len(payload) + truncated := false + if origSize > r.cfg.EventMaxPayload { + payload, truncated = truncateLargeFields(msg, r.cfg.EventTruncateField, r.cfg.EventMaxPayload) + } + + ev, err := r.repo.InsertEvent(ctx, &Event{ + SessionID: r.sessionID, + Direction: EventDirection(direction), + RPCKind: RPCKind(rpcKind), + Method: ptrIfNotEmpty(method), + Payload: payload, + PayloadSize: int32(origSize), + Truncated: truncated, + }) + if err != nil { + r.log.Error("acp.relay.persist_event", "session_id", r.sessionID, "err", err.Error()) + return nil + } + return &EventEnvelope{ + ID: ev.ID, + Direction: direction, + Kind: rpcKind, + Method: method, + Payload: payload, + Truncated: truncated, + } +} + +func classifyKind(m *Message) string { + switch { + case m.IsRequest(): + return string(RPCRequest) + case m.IsNotification(): + return string(RPCNotification) + case m.IsResponse(): + if m.Error != nil { + return string(RPCError) + } + return string(RPCResponse) + } + return "unknown" +} + +func ptrIfNotEmpty(s string) *string { + if s == "" { + return nil + } + return &s +} + +// truncateLargeFields replaces an oversize payload with a placeholder summary. +// Spec §3.5. MVP: full-payload truncate (don't try to be clever about which +// field is large — agents may put data in many places). +func truncateLargeFields(m *Message, fieldMax, totalMax int) ([]byte, bool) { + original, _ := json.Marshal(m) + if len(original) <= totalMax { + return original, false + } + hash := simpleSHA256Hex(original) + placeholder := fmt.Sprintf("%s... [TRUNCATED size=%d sha256=%s]", + string(original[:fieldMax]), len(original), hash) + out, _ := json.Marshal(map[string]any{ + "_truncated_payload": placeholder, + "_method": m.Method, + }) + return out, true +} + +func simpleSHA256Hex(b []byte) string { + h := sha256.Sum256(b) + return hex.EncodeToString(h[:]) +}