core/server/events/bus.go

package events

import (
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// Bus is a labeled publish/subscribe event bus.
// Handlers subscribe to patterns (exact or trailing wildcard).
// Publishing fans out to all matching subscribers.
type Bus struct {
	mu            sync.RWMutex
	subs          map[string][]*subscription
	seq           uint64 // subscription ID counter
	broadcastHook func(Event) // called after Publish for cross-pod fan-out; nil-safe
	publishCount  atomic.Int64
	deliverCount  atomic.Int64
}

type subscription struct {
	id      uint64
	pattern string
	handler Handler
}

// NewBus creates a new event bus.
func NewBus() *Bus {
	return &Bus{
		subs: make(map[string][]*subscription),
	}
}

// SetBroadcastHook registers a function called after every Publish.
// Used by the Postgres LISTEN/NOTIFY adapter to fan out across pods.
// The hook is NOT called by publishLocal (avoids re-broadcast loops).
func (b *Bus) SetBroadcastHook(fn func(Event)) {
	b.mu.Lock()
	b.broadcastHook = fn
	b.mu.Unlock()
}

// Subscribe registers a handler for a label pattern.
// Returns an unsubscribe function.
//
// Patterns:
//
//	"workflow.assigned.abc123"  — exact match
//	"workflow.assigned.*"      — wildcard: matches workflow.assigned.{anything}
//	"workflow.*"               — wildcard: matches workflow.{anything}
//	"*"                        — matches all events
func (b *Bus) Subscribe(pattern string, handler Handler) func() {
	b.mu.Lock()
	b.seq++
	sub := &subscription{id: b.seq, pattern: pattern, handler: handler}
	b.subs[pattern] = append(b.subs[pattern], sub)
	b.mu.Unlock()

	return func() {
		b.mu.Lock()
		defer b.mu.Unlock()
		subs := b.subs[pattern]
		for i, s := range subs {
			if s.id == sub.id {
				b.subs[pattern] = append(subs[:i], subs[i+1:]...)
				break
			}
		}
	}
}

// Publish dispatches an event to all matching subscribers.
// Handlers are called synchronously in subscription order.
// Use PublishAsync for non-blocking dispatch.
// After local dispatch, calls the broadcastHook if set (Postgres fan-out).
func (b *Bus) Publish(event Event) {
	b.publishLocal(event)

	b.mu.RLock()
	hook := b.broadcastHook
	b.mu.RUnlock()
	if hook != nil {
		hook(event)
	}
}

// publishLocal dispatches to local subscribers only, without triggering
// the broadcastHook. Used by the Postgres listener to re-publish remote
// events without causing an infinite re-broadcast loop.
func (b *Bus) publishLocal(event Event) {
	b.publishCount.Add(1)

	b.mu.RLock()
	var matched []Handler
	for pattern, subs := range b.subs {
		if match(event.Label, pattern) {
			for _, s := range subs {
				matched = append(matched, s.handler)
			}
		}
	}
	b.mu.RUnlock()

	for _, h := range matched {
		b.deliverCount.Add(1)
		h(event)
	}
}

// PublishAsync dispatches an event to all matching subscribers
// in separate goroutines. Useful for I/O-heavy handlers.
// Also calls the broadcastHook asynchronously if set.
func (b *Bus) PublishAsync(event Event) {
	b.publishCount.Add(1)

	b.mu.RLock()
	var matched []Handler
	for pattern, subs := range b.subs {
		if match(event.Label, pattern) {
			for _, s := range subs {
				matched = append(matched, s.handler)
			}
		}
	}
	hook := b.broadcastHook
	b.mu.RUnlock()

	for _, h := range matched {
		b.deliverCount.Add(1)
		go h(event)
	}
	if hook != nil {
		go hook(event)
	}
}

// PublishCount returns the cumulative number of events published.
func (b *Bus) PublishCount() int64 { return b.publishCount.Load() }

// DeliverCount returns the cumulative number of subscriber deliveries.
func (b *Bus) DeliverCount() int64 { return b.deliverCount.Load() }

// match checks if a concrete label matches a subscription pattern.
//
//	"chat.message.abc" matches "chat.message.abc"  (exact)
//	"chat.message.abc" matches "chat.message.*"    (wildcard)
//	"chat.message.abc" matches "chat.*"            (wildcard)
//	"chat.message.abc" matches "*"                 (global wildcard)
func match(label, pattern string) bool {
	if pattern == "*" {
		return true
	}
	if pattern == label {
		return true
	}
	if !strings.HasSuffix(pattern, "*") {
		return false
	}
	prefix := pattern[:len(pattern)-1] // "chat.message." from "chat.message.*"
	return strings.HasPrefix(label, prefix)
}

// WaitFor subscribes to an exact label and blocks until an event arrives
// or the timeout elapses. Returns the event and true, or a zero Event and
// false on timeout. The subscription is automatically removed after.
//
// Used by the browser tool bridge: the server publishes tool.call.{id}
// to the client and then calls WaitFor("tool.result.{id}", 30s) to block
// until the browser returns the result.
func (b *Bus) WaitFor(label string, timeout time.Duration) (Event, bool) {
	ch := make(chan Event, 1)
	unsub := b.Subscribe(label, func(e Event) {
		select {
		case ch <- e:
		default:
		}
	})
	defer unsub()

	select {
	case e := <-ch:
		return e, true
	case <-time.After(timeout):
		return Event{}, false
	}
}