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 (
	"encoding/json"
	"fmt"
	"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)
}

// 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
}