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 is an agent-to-agent communication endpoint bound to a single agent id.
// It is safe for concurrent use: the only mutable field is an atomic sequence
// counter used to mint unique ask ids, and every in-flight ask carries its own
// reply subscription and channel.
type Probe struct {
	bus        bus.Bus
	agentID    string
	seq        atomic.Uint64
	askTimeout time.Duration
	log        *slog.Logger
}

// Option customizes a Probe built by New.
type Option func(*Probe)

// WithLogger replaces the probe's logger. The default is a discard logger so the
// probe is quiet until configured.
func WithLogger(l *slog.Logger) Option {
	return func(p *Probe) { p.log = newLogger(l) }
}

// WithAskTimeout sets how long Ask waits for a reply when the caller's context
// carries no shorter deadline. A non-positive value keeps DefaultAskTimeout.
func WithAskTimeout(d time.Duration) Option {
	return func(p *Probe) { p.askTimeout = d }
}

// New builds a Probe bound to agentID. The agentID is sanitized for use in bus
// subjects (see token); an empty id panics because every probe must be routable.
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 returns the sanitized id this Probe is bound to.
func (p *Probe) AgentID() string { return p.agentID }

// Ask sends a question to to and blocks until a matching Reply arrives, an ask
// deadline elapses, or ctx is canceled. It returns the reply on success;
// otherwise it returns ErrTimeout when the deadline elapsed without a reply,
// ErrNoReply when the ask could not be delivered at all, or the wrapped context
// error when ctx was canceled.
//
// The caller's context supplies a deadline if one is set; otherwise
// p.askTimeout bounds the wait. The reply subscription is created before the
// ask is published, so a fast answerer cannot be missed.
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 publishes an answer to the ask identified by r.AskID, addressed to to.
// The answerer defaults to this Probe's agent id when r.From is empty, so an
// answering handler can simply pass the reply text and payload.
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 is a convenience for replying with plain text. It publishes a Reply
// for askID to the asking agent to.
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 delivers an event addressed specifically to to. The emitter defaults to
// this Probe's agent id when ev.From is empty.
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 broadcasts an event to every agent subscribed for events. The
// emitter defaults to this Probe's agent id when ev.From is empty.
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 mints a unique ask id from the agent id and an atomic counter. It is
// safe for concurrent calls.
func (p *Probe) nextAskID() string {
	return fmt.Sprintf("%s-%d", p.agentID, p.seq.Add(1))
}

// newLogger returns l tagged with the orchestrations/comm source, swapping a nil
// logger for a discard one so the probe can always log safely.
func newLogger(l *slog.Logger) *slog.Logger {
	if l == nil {
		l = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	return logger.WithSource(l, "orchestration/comm")
}