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") }