agentic-coding-workflow/internal/acp/relay.go

// relay.go 实现 ACP 双向 JSON-RPC 中继：
//   - reader: 子进程 stdout → 解析 JSON-RPC → 落 events → fanout / handler dispatch
//   - writer: 服务端发起 / WS 客户端透传 → 子进程 stdin
//   - subscribers: 多 WS tab 共享 reader 的 fanout
//
// id 命名空间（spec §5.2 / §7.2）：
//   - 服务端发起：int64（atomic 递增）
//   - 客户端 WS 发起：string（前端自生成）
//   - inflight 仅存 int64 → response chan；string id 的 response 当 event fanout
package acp

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"sync"
	"sync/atomic"

	"github.com/google/uuid"

	"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
)

// EventEnvelope 是推到 WS 的事件结构（与落库的 Event 字段同步，但带 ID 数值化方便 JSON）。
type EventEnvelope struct {
	ID        int64           `json:"id"`
	Direction string          `json:"direction"` // 'in' | 'out'
	Kind      string          `json:"kind"`      // 'request'|'response'|'notification'|'error'|'session_terminated'|'slow_consumer_disconnect'
	Method    string          `json:"method,omitempty"`
	Payload   json.RawMessage `json:"payload"`
	Truncated bool            `json:"truncated"`
}

// Subscriber 是单个 WS tab 的订阅句柄。
type Subscriber struct {
	ID     uuid.UUID
	UserID uuid.UUID
	Send   chan *EventEnvelope
	closed atomic.Bool
}

// Relay 是单 session 的中继器。每个活跃 session 一个 Relay 实例。
type Relay struct {
	sessionID uuid.UUID

	// 协议层
	enc *Encoder
	dec *Decoder

	// 服务端发起请求等响应（id=int64）
	inflight sync.Map // int64 → chan *Message
	nextID   atomic.Int64

	// WS subscribers（多 tab fanout）
	subsMu sync.Mutex
	subs   map[uuid.UUID]*Subscriber

	// 写入路径
	serverInitiated chan *Message // 服务端发起的 request/notification
	clientInitiated chan *Message // WS 上行（已校验 method 白名单）

	// 依赖
	repo        Repository
	fsHandler   *handlers.FsHandler
	permHandler *handlers.PermissionHandler
	log         *slog.Logger
	cfg         RelayConfig

	// 终止控制
	closed atomic.Bool
	done   chan struct{}
}

// RelayConfig 影响事件入库截断行为。
type RelayConfig struct {
	EventMaxPayload    int
	EventTruncateField int
	WSSendBuffer       int
}

// NewRelay 构造一个空 Relay；启动 goroutine 走 Run（D7 实现）。
func NewRelay(sessionID uuid.UUID, dec *Decoder, enc *Encoder,
	repo Repository, fs *handlers.FsHandler, perm *handlers.PermissionHandler,
	cfg RelayConfig, log *slog.Logger) *Relay {
	if log == nil {
		log = slog.Default()
	}
	return &Relay{
		sessionID:       sessionID,
		enc:             enc,
		dec:             dec,
		subs:            map[uuid.UUID]*Subscriber{},
		serverInitiated: make(chan *Message, 32),
		clientInitiated: make(chan *Message, 32),
		repo:            repo,
		fsHandler:       fs,
		permHandler:     perm,
		log:             log,
		cfg:             cfg,
		done:            make(chan struct{}),
	}
}

// Subscribe 注册一个 WS tab 订阅。
func (r *Relay) Subscribe(userID uuid.UUID) *Subscriber {
	sub := &Subscriber{
		ID:     uuid.New(),
		UserID: userID,
		Send:   make(chan *EventEnvelope, r.cfg.WSSendBuffer),
	}
	r.subsMu.Lock()
	r.subs[sub.ID] = sub
	r.subsMu.Unlock()
	return sub
}

// Unsubscribe 移除订阅；幂等。
func (r *Relay) Unsubscribe(subID uuid.UUID) {
	r.subsMu.Lock()
	defer r.subsMu.Unlock()
	if sub, ok := r.subs[subID]; ok {
		if sub.closed.CompareAndSwap(false, true) {
			close(sub.Send)
		}
		delete(r.subs, subID)
	}
}

// fanout 把一条 envelope 推给所有 subscribers。慢消费者按 §7.6 自动断开。
func (r *Relay) fanout(env *EventEnvelope) {
	r.subsMu.Lock()
	defer r.subsMu.Unlock()
	for id, sub := range r.subs {
		select {
		case sub.Send <- env:
		default:
			// buffer 满 → 慢消费者，标 closed + 通知 + close channel
			if sub.closed.CompareAndSwap(false, true) {
				disc := &EventEnvelope{Kind: "slow_consumer_disconnect"}
				select {
				case sub.Send <- disc:
				default:
				}
				close(sub.Send)
				delete(r.subs, id)
				r.log.Warn("acp.relay.slow_subscriber_dropped",
					"session_id", r.sessionID, "sub_id", id)
			}
		}
	}
}

// Close 优雅关闭：先广播 final event，再关 channel，再标 closed。
// supervisor.onExit 调用，确保所有 WS 看到 session_terminated。
func (r *Relay) Close(status string, exitCode *int32) {
	if !r.closed.CompareAndSwap(false, true) {
		return
	}
	final := &EventEnvelope{
		Kind: "session_terminated",
		Payload: mustJSON(map[string]any{
			"status":    status,
			"exit_code": exitCode,
		}),
	}
	r.fanout(final)

	r.subsMu.Lock()
	for id, sub := range r.subs {
		if sub.closed.CompareAndSwap(false, true) {
			close(sub.Send)
		}
		delete(r.subs, id)
	}
	r.subsMu.Unlock()

	close(r.done)
}

// FanoutControl 把一条控制类 envelope（如 permission_request/permission_resolved）
// 推给所有 WS 订阅者。这类帧不落 acp_events、ID 恒为 0，由 PermissionService 调用。
// relay 已关闭时静默丢弃，避免向已关 channel 写入。
func (r *Relay) FanoutControl(env *EventEnvelope) {
	if r.closed.Load() {
		return
	}
	r.fanout(env)
}

// Done returns a channel closed when Relay is terminated.
func (r *Relay) Done() <-chan struct{} { return r.done }

func mustJSON(v any) json.RawMessage {
	b, err := json.Marshal(v)
	if err != nil {
		return json.RawMessage(fmt.Sprintf(`{"_marshal_err":%q}`, err.Error()))
	}
	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.
// 在独立 goroutine 中运行，故 defer recover 兜底任何 handler panic：记录日志并回送
// JSON-RPC internal error，避免 agent 构造的请求 panic 掉整个进程（DoS）。
func (r *Relay) handleAgentRequest(ctx context.Context, sess handlers.SessionContext, req *Message) {
	defer func() {
		if rec := recover(); rec != nil {
			r.log.Error("acp.relay.handle_request_panic",
				"session_id", r.sessionID, "method", req.Method, "panic", fmt.Sprintf("%v", rec))
			resp := NewResponseErr(req.ID, JSONRPCInternalError, "internal error")
			select {
			case r.serverInitiated <- resp:
			case <-ctx.Done():
			}
		}
	}()

	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, string(req.ID), 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[:])
}