diff --git a/internal/runtime/orchestration/comm/comm.go b/internal/runtime/orchestration/comm/comm.go new file mode 100644 index 0000000..cfc02c4 --- /dev/null +++ b/internal/runtime/orchestration/comm/comm.go @@ -0,0 +1,178 @@ +// Package comm implements agent-to-agent communication for the Cyrene v2 +// orchestration layer (DevDocs/design/05-子代理.md §4). Subagents and the main +// agent exchange messages over the in-process bus (internal/bus) using a +// request/answer protocol built on three message kinds: ask, reply and event. +// +// # Protocol +// +// A Probe is bound to one agent id. It can ask another agent a question and +// wait for its reply (ask/reply), or fire events at another agent or to anyone +// listening (event). Messages are published to the bus with dot-separated +// subjects so the bus's NATS-like wildcard matching routes them to the right +// agent; see the "Subject convention" section below. +// +// # Subject convention +// +// Subjects are built from sanitized agent ids (any rune outside [A-Za-z0-9_-] +// is replaced by "_", so an id never breaks a subject token): +// +// ask: cyrene.agent...ask +// reply: cyrene.agent...reply. +// event (to): cyrene.agent...event +// event (bc): cyrene.agent..event +// +// The reply subject carries the ask id as its last token, giving each in-flight +// ask a private reply address so concurrent asks from the same agent never +// cross replies. To receive asks addressed to it, an agent subscribes to +// cyrene.agent.*..ask. To receive events (both broadcasts and events +// addressed to it) it subscribes to cyrene.agent.*.event and to +// cyrene.agent.*..event. +// +// # Delivery semantics +// +// Ask blocks until a reply arrives, an ask-timeout elapses, or the passed +// context is canceled; it returns ErrTimeout on the former and the wrapped +// context error on the latter. Reply and the event senders are fire-and-forget +// publishes that return the bus error (if any). Handler callbacks are invoked +// on the bus's per-subscription worker goroutines, so an EventHandler may be +// called concurrently from two subscriptions and must be safe for concurrent +// use; an AskHandler for one agent is serialized by the bus. +package comm + +import ( + "errors" + "strings" + "time" +) + +// Kind is the wire kind of an agent-to-agent message. +type Kind string + +const ( + // KindAsk is a question from one agent to another, answered by a reply. + KindAsk Kind = "ask" + // KindReply is the answer to a previously published ask. + KindReply Kind = "reply" + // KindEvent is a one-directional notification, broadcast or addressed. + KindEvent Kind = "event" +) + +// subjectPrefix is the fixed first two tokens of every agent-to-agent subject. +const subjectPrefix = "cyrene.agent" + +// DefaultAskTimeout caps how long Ask waits for a reply when the caller's +// context carries no shorter deadline. It bounds a request/answer round so a +// silent agent never hangs its caller. +const DefaultAskTimeout = 30 * time.Second + +// Sentinel errors returned by the package. Use errors.Is to test for them; they +// are returned directly or wrapped with %w so the underlying cause stays +// available. +var ( + // ErrTimeout is returned by Ask when no reply arrives before the ask + // deadline (the caller's context deadline or DefaultAskTimeout) elapses. + ErrTimeout = errors.New("comm: reply timeout") + + // ErrNoReply is returned by Ask when the ask cannot be delivered (the bus + // rejected the publish/reply subscription), so no reply is expected. + ErrNoReply = errors.New("comm: no reply") +) + +// Ask is a question sent from one agent to another. +type Ask struct { + // ID uniquely identifies this ask within the answering exchange. + ID string + // From is the id of the asking agent (the sender). + From string + // To is the id of the agent that should answer (the receiver). + To string + // Question is the human-readable question text. + Question string + // Data carries an optional structured payload alongside the question. + Data any +} + +// Reply is the answer to a previously published ask. +type Reply struct { + // AskID is the ID of the ask this reply answers. + AskID string + // From is the id of the replying agent (the sender). + From string + // To is the id of the asking agent (the receiver). + To string + // Text is the human-readable answer text. + Text string + // Data carries an optional structured payload alongside the answer. + Data any +} + +// Event is a one-directional notification, either broadcast or addressed to one +// agent. It has no reply. +type Event struct { + // From is the id of the emitting agent. + From string + // Name is a short event name, e.g. "task.progress" or "tool.used". + Name string + // Data carries an optional structured payload. + Data any +} + +// token sanitizes an agent id so it can be used as a single subject token: any +// rune outside [A-Za-z0-9_-] is replaced with "_". This keeps the id useful for +// routing while guaranteeing dots (and the other subject metacharacters) never +// leak into a token. +func token(s string) string { + var b strings.Builder + for _, r := range s { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_': + b.WriteRune(r) + default: + b.WriteRune('_') + } + } + return b.String() +} + +// subjectFor builds the base subject for a from/to message of the given kind. +func subjectFor(from, to string, kind Kind) string { + return subjectPrefix + "." + token(from) + "." + token(to) + "." + string(kind) +} + +// askSubject is the subject an ask from→to is published to. +func askSubject(from, to string) string { + return subjectFor(from, to, KindAsk) +} + +// replySubject is the private subject a reply from→to answering askID is +// published to. The trailing askID token gives each in-flight ask its own +// address so replies never cross between concurrent asks. +func replySubject(from, to, askID string) string { + return subjectPrefix + "." + token(from) + "." + token(to) + "." + string(KindReply) + "." + token(askID) +} + +// eventSubject is the subject of an event addressed from→to. +func eventSubject(from, to string) string { + return subjectFor(from, to, KindEvent) +} + +// broadcastEventSubject is the subject of a broadcast event emitted by from. +func broadcastEventSubject(from string) string { + return subjectPrefix + "." + token(from) + "." + string(KindEvent) +} + +// askReceivePattern is the subscription pattern an agent named me uses to +// receive asks addressed to it from any sender. +func askReceivePattern(me string) string { + return subjectPrefix + ".*." + token(me) + "." + string(KindAsk) +} + +// eventReceivePatterns are the subscription patterns an agent named me uses to +// receive every event it should see: broadcasts from any sender and events +// addressed specifically to it. +func eventReceivePatterns(me string) []string { + return []string{ + subjectPrefix + ".*." + string(KindEvent), + subjectPrefix + ".*." + token(me) + "." + string(KindEvent), + } +} diff --git a/internal/runtime/orchestration/comm/comm_test.go b/internal/runtime/orchestration/comm/comm_test.go new file mode 100644 index 0000000..26644ae --- /dev/null +++ b/internal/runtime/orchestration/comm/comm_test.go @@ -0,0 +1,359 @@ +package comm + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/bus" +) + +// recvEvent blocks until an event arrives on ch or d elapses, failing the test +// on timeout. It makes delivery assertions deterministic instead of relying on a +// bare sleep. +func recvEvent(t *testing.T, ch <-chan Event, d time.Duration) Event { + t.Helper() + select { + case e := <-ch: + return e + case <-time.After(d): + t.Fatalf("timed out waiting for an event after %v", d) + return Event{} + } +} + +func TestSubjectConvention(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + got string + want string + }{ + {"ask", askSubject("alice", "bob"), "cyrene.agent.alice.bob.ask"}, + {"reply carries ask id", replySubject("bob", "alice", "alice-3"), "cyrene.agent.bob.alice.reply.alice-3"}, + {"event addressed", eventSubject("alice", "bob"), "cyrene.agent.alice.bob.event"}, + {"event broadcast", broadcastEventSubject("alice"), "cyrene.agent.alice.event"}, + {"ask receive pattern", askReceivePattern("bob"), "cyrene.agent.*.bob.ask"}, + {"agent ids sanitized", askSubject("a.b", "c d"), "cyrene.agent.a_b.c_d.ask"}, + } + for _, c := range cases { + c := c + t.Run(c.name, func(t *testing.T) { + t.Parallel() + if c.got != c.want { + t.Fatalf("subject = %q, want %q", c.got, c.want) + } + }) + } +} + +func TestEventReceivePatterns(t *testing.T) { + t.Parallel() + + got := eventReceivePatterns("bob") + want := []string{"cyrene.agent.*.event", "cyrene.agent.*.bob.event"} + if len(got) != len(want) { + t.Fatalf("patterns = %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("patterns[%d] = %q, want %q", i, got[i], want[i]) + } + } +} + +func TestAskReply_RoundTrip(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + bob := New(b, "bob") + + unsub, err := bob.SubscribeAsk(context.Background(), func(ctx context.Context, a Ask) error { + return bob.Answer(ctx, a.From, a.ID, "pong:"+a.Question, a.Data) + }) + if err != nil { + t.Fatalf("SubscribeAsk: %v", err) + } + defer unsub() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + reply, err := alice.Ask(ctx, "bob", "ping", "payload") + if err != nil { + t.Fatalf("Ask: %v", err) + } + if reply.Text != "pong:ping" { + t.Fatalf("reply.Text = %q, want pong:ping", reply.Text) + } + if reply.From != "bob" { + t.Fatalf("reply.From = %q, want bob", reply.From) + } + if reply.To != "alice" { + t.Fatalf("reply.To = %q, want alice", reply.To) + } + if reply.Data != "payload" { + t.Fatalf("reply.Data = %#v, want payload", reply.Data) + } +} + +func TestAsk_TimeoutNoReply(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + // bob has no ask handler, so nothing ever replies. + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + _, err := alice.Ask(ctx, "bob", "hello", nil) + if !errors.Is(err, ErrTimeout) { + t.Fatalf("Ask err = %v, want ErrTimeout", err) + } +} + +func TestAsk_ContextCanceled(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + + ctx, cancel := context.WithCancel(context.Background()) + got := make(chan error, 1) + go func() { + _, err := alice.Ask(ctx, "bob", "hello", nil) + got <- err + }() + time.Sleep(50 * time.Millisecond) + cancel() + + select { + case err := <-got: + if !errors.Is(err, context.Canceled) { + t.Fatalf("Ask err = %v, want context.Canceled", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Ask did not return after ctx cancel") + } +} + +func TestAsk_ContextPreCanceled(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := alice.Ask(ctx, "bob", "q", nil) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Ask err = %v, want context.Canceled", err) + } +} + +func TestEventBroadcast(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + bob := New(b, "bob") + + events := make(chan Event, 4) + unsub, err := bob.SubscribeEvent(context.Background(), func(ctx context.Context, e Event) error { + events <- e + return nil + }) + if err != nil { + t.Fatalf("SubscribeEvent: %v", err) + } + defer unsub() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + if err := alice.PublishEvent(ctx, Event{Name: "tick", Data: 7}); err != nil { + t.Fatalf("PublishEvent: %v", err) + } + e := recvEvent(t, events, 2*time.Second) + if e.Name != "tick" || e.Data != 7 || e.From != "alice" { + t.Fatalf("event = %+v, want name=tick data=7 from=alice", e) + } +} + +func TestEventSend_AddressedAndIsolated(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + bob := New(b, "bob") + carol := New(b, "carol") + + bobEvents := make(chan Event, 4) + carolEvents := make(chan Event, 4) + + unsubBob, err := bob.SubscribeEvent(context.Background(), func(ctx context.Context, e Event) error { + bobEvents <- e + return nil + }) + if err != nil { + t.Fatalf("SubscribeEvent bob: %v", err) + } + defer unsubBob() + unsubCarol, err := carol.SubscribeEvent(context.Background(), func(ctx context.Context, e Event) error { + carolEvents <- e + return nil + }) + if err != nil { + t.Fatalf("SubscribeEvent carol: %v", err) + } + defer unsubCarol() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + if err := alice.Send(ctx, "bob", Event{Name: "direct"}); err != nil { + t.Fatalf("Send: %v", err) + } + e := recvEvent(t, bobEvents, 2*time.Second) + if e.Name != "direct" || e.From != "alice" { + t.Fatalf("bob event = %+v, want name=direct from=alice", e) + } + + // An event addressed to bob must not reach carol. + select { + case e := <-carolEvents: + t.Fatalf("carol unexpectedly received %+v", e) + case <-time.After(100 * time.Millisecond): + } +} + +func TestEventHandler_Concurrent(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + emitter := New(b, "emit") + recv := New(b, "recv") + + const n = 100 + var got atomic.Int64 + done := make(chan struct{}) + unsub, err := recv.SubscribeEvent(context.Background(), func(ctx context.Context, e Event) error { + if got.Add(1) == n { + close(done) + } + return nil + }) + if err != nil { + t.Fatalf("SubscribeEvent: %v", err) + } + defer unsub() + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + // Fire a mix of broadcasts and events addressed to recv so both event + // subscriptions (and their separate worker goroutines) invoke the handler + // concurrently. + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + if i%2 == 0 { + _ = emitter.PublishEvent(ctx, Event{Name: "bc", Data: i}) + } else { + _ = emitter.Send(ctx, "recv", Event{Name: "to", Data: i}) + } + }(i) + } + wg.Wait() + + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatalf("only %d of %d events delivered", got.Load(), n) + } +} + +func TestAsk_ConcurrentDistinct(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + bob := New(b, "bob") + + unsub, err := bob.SubscribeAsk(context.Background(), func(ctx context.Context, a Ask) error { + return bob.Answer(ctx, a.From, a.ID, "ack:"+a.Question, a.Data) + }) + if err != nil { + t.Fatalf("SubscribeAsk: %v", err) + } + defer unsub() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + const n = 20 + var wg sync.WaitGroup + errs := make(chan error, n) + for i := 0; i < n; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + q := fmt.Sprintf("q%d", i) + reply, err := alice.Ask(ctx, "bob", q, i) + if err != nil { + errs <- fmt.Errorf("ask %d: %w", i, err) + return + } + if reply.Text != "ack:"+q { + errs <- fmt.Errorf("ask %d: reply.Text = %q, want %q", i, reply.Text, "ack:"+q) + return + } + if reply.Data != i { + errs <- fmt.Errorf("ask %d: reply.Data = %#v, want %d", i, reply.Data, i) + return + } + }(i) + } + wg.Wait() + close(errs) + for err := range errs { + t.Fatal(err) + } +} + +func TestSend_NoErrorAndNoReceiver(t *testing.T) { + t.Parallel() + + b := bus.New() + defer b.Close() + alice := New(b, "alice") + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + // Publishing to a subject with no subscriber is a no-op, not an error. + if err := alice.Send(ctx, "ghost", Event{Name: "nobody"}); err != nil { + t.Fatalf("Send to lost receiver: %v", err) + } + if err := alice.PublishEvent(ctx, Event{Name: "nobody"}); err != nil { + t.Fatalf("PublishEvent with no subscriber: %v", err) + } +} diff --git a/internal/runtime/orchestration/comm/probe.go b/internal/runtime/orchestration/comm/probe.go new file mode 100644 index 0000000..62d2716 --- /dev/null +++ b/internal/runtime/orchestration/comm/probe.go @@ -0,0 +1,182 @@ +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") +} diff --git a/internal/runtime/orchestration/comm/subscribe.go b/internal/runtime/orchestration/comm/subscribe.go new file mode 100644 index 0000000..ab3c197 --- /dev/null +++ b/internal/runtime/orchestration/comm/subscribe.go @@ -0,0 +1,95 @@ +package comm + +import ( + "context" + "fmt" + "sync" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/bus" +) + +// AskHandler processes an Ask addressed to this agent. It runs on the bus's +// per-subscription worker goroutine; for a single agent asks are serialized. A +// handler should answer via (p *Probe).Reply or Answer; a non-nil returned +// error is logged by the probe and discarded. +type AskHandler func(ctx context.Context, a Ask) error + +// EventHandler processes an Event addressed to, or broadcast to, this agent. +// Because an agent subscribes to two event patterns (broadcast and addressed), +// the bus may invoke a handler concurrently from two worker goroutines, so an +// EventHandler must be safe for concurrent use. +type EventHandler func(ctx context.Context, e Event) error + +// SubscribeAsk registers h to receive every Ask addressed to this agent and +// returns an unsubscribe function that removes the registration. When ctx is +// canceled the registration is removed automatically; the returned function is +// idempotent. +func (p *Probe) SubscribeAsk(ctx context.Context, h AskHandler) (func(), error) { + if h == nil { + return nil, fmt.Errorf("comm: subscribe ask: handler must not be nil") + } + pattern := askReceivePattern(p.agentID) + unsub, err := p.bus.Subscribe(ctx, pattern, func(bctx context.Context, e bus.Event) error { + a, ok := e.Payload.(Ask) + if !ok { + return fmt.Errorf("comm: ask payload type %T, want Ask", e.Payload) + } + if err := h(bctx, a); err != nil { + p.log.Warn("ask handler failed", "from", a.From, "ask_id", a.ID, "err", err) + } + return nil + }) + if err != nil { + return nil, fmt.Errorf("comm: subscribe ask: %w", err) + } + return func() { unsub() }, nil +} + +// SubscribeEvent registers h to receive both broadcast events and events +// addressed to this agent. It returns an unsubscribe function that removes both +// registrations; it is idempotent and safe to call more than once. When ctx is +// canceled both registrations are removed automatically. +// +// The handler reaches the caller from two subscription worker goroutines and +// may also run concurrently with a SubscribeAsk handler, so it must be safe for +// concurrent use. +func (p *Probe) SubscribeEvent(ctx context.Context, h EventHandler) (func(), error) { + if h == nil { + return nil, fmt.Errorf("comm: subscribe event: handler must not be nil") + } + + wrap := func(bctx context.Context, e bus.Event) error { + ev, ok := e.Payload.(Event) + if !ok { + return fmt.Errorf("comm: event payload type %T, want Event", e.Payload) + } + if err := h(bctx, ev); err != nil { + p.log.Warn("event handler failed", "from", ev.From, "name", ev.Name, "err", err) + } + return nil + } + + var ( + mu sync.Mutex + unsubs []func() + ) + combined := func() { + mu.Lock() + defer mu.Unlock() + for _, u := range unsubs { + u() + } + } + + for _, pattern := range eventReceivePatterns(p.agentID) { + u, err := p.bus.Subscribe(ctx, pattern, wrap) + if err != nil { + combined() + return nil, fmt.Errorf("comm: subscribe event %q: %w", pattern, err) + } + mu.Lock() + unsubs = append(unsubs, u) + mu.Unlock() + } + return combined, nil +}