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

package run

import (
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
)

// LogEnvelope 是推到 WS / 经日志 tail 返回的一行日志或一个控制帧。
//   - kind="log"：普通日志行（level=info 来自 stdout，level=error 来自 stderr）
//   - kind="state"：运行态变化（Status/ExitCode 有意义）
//   - kind="slow_consumer_disconnect"：慢消费者被动断开的控制帧
type LogEnvelope struct {
	ID       int64     `json:"id"`
	Kind     string    `json:"kind"`
	Level    string    `json:"level,omitempty"`
	Text     string    `json:"text,omitempty"`
	TS       time.Time `json:"ts"`
	Status   string    `json:"status,omitempty"`
	ExitCode *int32    `json:"exit_code,omitempty"`
}

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

// outRing 是固定行数的环形缓冲：满则淘汰最旧行。承载某次运行的合并 stdout+stderr
// 日志，既供 logs(tail/since) 查询，也供 onExit 取尾部写 last_error。
type outRing struct {
	mu    sync.Mutex
	lines []*LogEnvelope
	cap   int
}

func newOutRing(capacity int) *outRing {
	if capacity <= 0 {
		capacity = 1000
	}
	return &outRing{cap: capacity}
}

func (r *outRing) append(e *LogEnvelope) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.lines) >= r.cap {
		r.lines = r.lines[1:]
	}
	r.lines = append(r.lines, e)
}

// since 返回 ID > sinceID 的日志行（最多 limit 条；limit<=0 表示不限）。
func (r *outRing) since(sinceID int64, limit int) []*LogEnvelope {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]*LogEnvelope, 0, len(r.lines))
	for _, e := range r.lines {
		if e.ID > sinceID {
			out = append(out, e)
		}
	}
	if limit > 0 && len(out) > limit {
		out = out[len(out)-limit:]
	}
	return out
}

// tailText 反向重建尾部文本（不超过 maxBytes），供 onExit 写 last_error。
func (r *outRing) tailText(maxBytes int) string {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := ""
	for i := len(r.lines) - 1; i >= 0 && len(out) < maxBytes; i-- {
		out = r.lines[i].Text + "\n" + out
	}
	if len(out) > maxBytes {
		out = out[len(out)-maxBytes:]
	}
	return out
}

// runRelay 是单次运行（一个 RunProc）的日志中继：合并 stdout+stderr 进环形
// 缓冲并向所有 WS 订阅者 fanout。结构参照 acp/relay.go，但只做单向日志广播，
// 无 JSON-RPC 双向中继。
type runRelay struct {
	profileID uuid.UUID
	ring      *outRing
	nextID    atomic.Int64

	subsMu  sync.Mutex
	subs    map[uuid.UUID]*Subscriber
	sendBuf int

	closed atomic.Bool
	done   chan struct{}
	log    *slog.Logger
}

func newRunRelay(profileID uuid.UUID, bufLines, sendBuf int, log *slog.Logger) *runRelay {
	if log == nil {
		log = slog.Default()
	}
	if sendBuf <= 0 {
		sendBuf = 256
	}
	return &runRelay{
		profileID: profileID,
		ring:      newOutRing(bufLines),
		subs:      map[uuid.UUID]*Subscriber{},
		sendBuf:   sendBuf,
		done:      make(chan struct{}),
		log:       log,
	}
}

// emit 记录一行日志（落环形缓冲 + fanout）。drainPipe goroutine 调用。
func (r *runRelay) emit(level, text string) {
	env := &LogEnvelope{
		ID:    r.nextID.Add(1),
		Kind:  "log",
		Level: level,
		Text:  text,
		TS:    time.Now(),
	}
	r.ring.append(env)
	r.fanout(env)
}

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

// Unsubscribe 移除订阅；幂等。
func (r *runRelay) 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 推给所有订阅者；慢消费者（Send 满）自动断开，
// 不阻塞其他订阅者（同 acp relay 约定）。
func (r *runRelay) fanout(env *LogEnvelope) {
	r.subsMu.Lock()
	defer r.subsMu.Unlock()
	for id, sub := range r.subs {
		select {
		case sub.Send <- env:
		default:
			if sub.closed.CompareAndSwap(false, true) {
				select {
				case sub.Send <- &LogEnvelope{Kind: "slow_consumer_disconnect", TS: time.Now()}:
				default:
				}
				close(sub.Send)
				delete(r.subs, id)
				r.log.Warn("run.relay.slow_subscriber_dropped",
					"profile_id", r.profileID, "sub_id", id)
			}
		}
	}
}

// since 返回 ID > sinceID 的历史日志行（供 logs tail / WS 续传）。
func (r *runRelay) since(sinceID int64, limit int) []*LogEnvelope {
	return r.ring.since(sinceID, limit)
}

// tailText 返回尾部文本（供 onExit 写 last_error）。
func (r *runRelay) tailText(maxBytes int) string {
	return r.ring.tailText(maxBytes)
}

// Close 广播终止 state 帧后关闭所有订阅。幂等。
func (r *runRelay) Close(status string, exitCode *int32) {
	if !r.closed.CompareAndSwap(false, true) {
		return
	}
	r.fanout(&LogEnvelope{
		ID:       r.nextID.Add(1),
		Kind:     "state",
		Status:   status,
		ExitCode: exitCode,
		TS:       time.Now(),
	})
	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)
}