You've already forked agentic-coding-workflow
feat(acp): relay reader/writer goroutines + Call + persist with truncation
This commit is contained in:
@@ -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[:])
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user