Cyrene_v2/internal/runtime/orchestration/comm/probe.go

package comm

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"sync/atomic"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// Probe 是绑定到单个 agent id 的 agent 间通信端点。它可安全并发使用：唯一可变字段是
// 用于生成唯一 ask id 的原子序号计数器，且每个在途 ask 都带有自己的回复订阅与 channel。
type Probe struct {
	bus        bus.Bus
	agentID    string
	seq        atomic.Uint64
	askTimeout time.Duration
	log        *slog.Logger
}

// Option 定制由 New 构建的 Probe。
type Option func(*Probe)

// WithLogger 替换探针的 logger。默认是丢弃式 logger，因此在配置之前探针是安静的。
func WithLogger(l *slog.Logger) Option {
	return func(p *Probe) { p.log = newLogger(l) }
}

// WithAskTimeout 设置调用方上下文未携带更短期限时 Ask 等待回复的时长。非正值保持
// DefaultAskTimeout。
func WithAskTimeout(d time.Duration) Option {
	return func(p *Probe) { p.askTimeout = d }
}

// New 构建一个绑定到 agentID 的 Probe。agentID 会被净化以便用于总线 subject
// （见 token）；空 id 会 panic，因为每个探针都必须可路由。
func New(b bus.Bus, agentID string, opts ...Option) *Probe {
	if agentID == "" {
		panic("comm: New: agentID must not be empty")
	}
	p := &Probe{
		bus:        b,
		agentID:    token(agentID),
		askTimeout: DefaultAskTimeout,
		log:        newLogger(nil),
	}
	for _, opt := range opts {
		opt(p)
	}
	if p.askTimeout <= 0 {
		p.askTimeout = DefaultAskTimeout
	}
	return p
}

// AgentID 返回该 Probe 所绑定到的已净化 id。
func (p *Probe) AgentID() string { return p.agentID }

// Ask 向 to 提一个问题并阻塞到匹配的 Reply 到达、ask 期限到期，或 ctx 被取消。成功时
// 返回回复；否则期限到期而无回复时返回 ErrTimeout，ask 完全无法投递时返回 ErrNoReply，
// ctx 被取消时返回包装后的上下文错误。
//
// 调用方上下文设了期限就用它；否则由 p.askTimeout 限制等待时长。回复订阅先于 ask 发布
// 创建，因此再快的应答方也不会被错过。
func (p *Probe) Ask(ctx context.Context, to, question string, data any) (Reply, error) {
	if err := ctx.Err(); err != nil {
		return Reply{}, fmt.Errorf("comm: ask to %q: %w", to, err)
	}
	askID := p.nextAskID()
	from := p.agentID

	replyCh := make(chan Reply, 1)
	replySub := replySubject(to, from, askID)
	unsub, err := p.bus.Subscribe(ctx, replySub, func(bctx context.Context, e bus.Event) error {
		r, ok := e.Payload.(Reply)
		if !ok {
			return fmt.Errorf("comm: reply payload type %T, want Reply", e.Payload)
		}
		if r.AskID != askID {
			return fmt.Errorf("comm: reply ask id %q, want %q", r.AskID, askID)
		}
		select {
		case replyCh <- r:
		default:
		}
		return nil
	})
	if err != nil {
		return Reply{}, fmt.Errorf("comm: ask to %q: subscribe reply: %w: %w", to, ErrNoReply, err)
	}
	defer unsub()

	ask := Ask{ID: askID, From: from, To: to, Question: question, Data: data}
	if err := p.bus.Publish(ctx, askSubject(from, to), ask); err != nil {
		return Reply{}, fmt.Errorf("comm: ask to %q: publish ask: %w: %w", to, ErrNoReply, err)
	}

	waitCtx, cancel := context.WithTimeout(ctx, p.askTimeout)
	defer cancel()

	select {
	case r := <-replyCh:
		return r, nil
	case <-waitCtx.Done():
		if errors.Is(waitCtx.Err(), context.Canceled) {
			return Reply{}, fmt.Errorf("comm: ask to %q: %w", to, ctx.Err())
		}
		return Reply{}, fmt.Errorf("comm: ask to %q: %w", to, ErrTimeout)
	}
}

// Reply 向 to 发布对 r.AskID 所标识 ask 的应答。r.From 为空时，应答方默认为本 Probe
// 的 agent id，因此应答处理函数只需传入回复文本与载荷。
func (p *Probe) Reply(ctx context.Context, to string, r Reply) error {
	if r.From == "" {
		r.From = p.agentID
	}
	if r.To == "" {
		r.To = to
	}
	if err := p.bus.Publish(ctx, replySubject(r.From, to, r.AskID), r); err != nil {
		return fmt.Errorf("comm: reply to %q: %w", to, err)
	}
	return nil
}

// Answer 是以纯文本回复的便捷方式。它向提问 agent to 发布一条针对 askID 的 Reply。
func (p *Probe) Answer(ctx context.Context, to, askID, text string, data any) error {
	return p.Reply(ctx, to, Reply{AskID: askID, From: p.agentID, To: to, Text: text, Data: data})
}

// Send 投递一个专门发给 to 的事件。ev.From 为空时，发射方默认为本 Probe 的 agent id。
func (p *Probe) Send(ctx context.Context, to string, ev Event) error {
	if ev.From == "" {
		ev.From = p.agentID
	}
	if err := p.bus.Publish(ctx, eventSubject(ev.From, to), ev); err != nil {
		return fmt.Errorf("comm: send event to %q: %w", to, err)
	}
	return nil
}

// PublishEvent 向每个订阅了事件的 agent 广播一个事件。ev.From 为空时，发射方默认为本
// Probe 的 agent id。
func (p *Probe) PublishEvent(ctx context.Context, ev Event) error {
	if ev.From == "" {
		ev.From = p.agentID
	}
	if err := p.bus.Publish(ctx, broadcastEventSubject(ev.From), ev); err != nil {
		return fmt.Errorf("comm: publish event: %w", err)
	}
	return nil
}

// nextAskID 由 agent id 与一个原子计数器生成唯一的 ask id。可安全并发调用。
func (p *Probe) nextAskID() string {
	return fmt.Sprintf("%s-%d", p.agentID, p.seq.Add(1))
}

// newLogger 返回标注了来源 orchestration/comm 的 l，并把 nil logger 换成丢弃式
// logger，使探针总能安全记录日志。
func newLogger(l *slog.Logger) *slog.Logger {
	if l == nil {
		l = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	return logger.WithSource(l, "orchestration/comm")
}