unibots/shell/bus/bus.go

// Package bus provides in-process agent-to-agent message passing.
package bus

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"time"
)

// Well-known message kinds used by the orchestrator.
const (
	KindTask       = "task"        // orchestrator → bot: handle this question
	KindTaskResult = "task_result" // bot → orchestrator: here is my answer
)

// AgentID identifies an agent.
type AgentID string

// AgentMessage is a message between agents.
type AgentMessage struct {
	From    AgentID
	To      AgentID
	Kind    string
	Payload map[string]string
}

// Bus manages channels for inter-agent communication.
type Bus struct {
	mu       sync.RWMutex
	channels map[AgentID]chan AgentMessage

	replyMu  sync.Mutex
	replyChs map[string]chan AgentMessage // taskID → one-shot reply channel

	logger *slog.Logger
}

// New creates a new Bus.
func New(logger *slog.Logger) *Bus {
	return &Bus{
		channels: make(map[AgentID]chan AgentMessage),
		replyChs: make(map[string]chan AgentMessage),
		logger:   logger.With("component", "bus"),
	}
}

// Subscribe registers an agent and returns its receive channel.
func (b *Bus) Subscribe(id AgentID) <-chan AgentMessage {
	b.mu.Lock()
	defer b.mu.Unlock()
	ch := make(chan AgentMessage, 64)
	b.channels[id] = ch
	b.logger.Info("bus_subscribe", "agent", id)
	return ch
}

// Send delivers a message to an agent's channel.
func (b *Bus) Send(msg AgentMessage) error {
	b.mu.RLock()
	ch, ok := b.channels[msg.To]
	b.mu.RUnlock()
	if !ok {
		b.logger.Warn("bus_not_found", "to", msg.To, "from", msg.From, "kind", msg.Kind)
		return fmt.Errorf("agent %q not registered on bus", msg.To)
	}
	select {
	case ch <- msg:
		b.logger.Debug("bus_send", "from", msg.From, "to", msg.To, "kind", msg.Kind)
		return nil
	default:
		b.logger.Warn("bus_queue_full", "to", msg.To, "from", msg.From, "kind", msg.Kind)
		return fmt.Errorf("agent %q message queue full", msg.To)
	}
}

// SendAndWait sends a task message and blocks until a reply with the matching
// taskID arrives or the context expires. The caller must ensure the reply is
// routed via Reply().
func (b *Bus) SendAndWait(ctx context.Context, msg AgentMessage, taskID string, timeout time.Duration) (AgentMessage, error) {
	ch := make(chan AgentMessage, 1)
	b.replyMu.Lock()
	b.replyChs[taskID] = ch
	b.replyMu.Unlock()

	defer func() {
		b.replyMu.Lock()
		delete(b.replyChs, taskID)
		b.replyMu.Unlock()
	}()

	if err := b.Send(msg); err != nil {
		return AgentMessage{}, err
	}

	b.logger.Debug("bus_send_and_wait", "task", taskID, "to", msg.To, "timeout", timeout)

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	select {
	case reply := <-ch:
		return reply, nil
	case <-timer.C:
		b.logger.Warn("bus_timeout", "task", taskID, "to", msg.To, "timeout", timeout)
		return AgentMessage{}, fmt.Errorf("task %s: delegation timeout after %s", taskID, timeout)
	case <-ctx.Done():
		return AgentMessage{}, ctx.Err()
	}
}

// Reply routes a task_result message to the waiting SendAndWait caller.
// If no one is waiting for this taskID, it falls back to regular Send.
func (b *Bus) Reply(taskID string, msg AgentMessage) error {
	b.replyMu.Lock()
	ch, ok := b.replyChs[taskID]
	b.replyMu.Unlock()

	if ok {
		select {
		case ch <- msg:
			b.logger.Debug("bus_reply", "task", taskID, "from", msg.From)
			return nil
		default:
			b.logger.Warn("bus_reply_full", "task", taskID)
			return fmt.Errorf("reply channel full for task %s", taskID)
		}
	}
	// Fallback: deliver via regular channel
	return b.Send(msg)
}

// Unsubscribe removes an agent from the bus.
func (b *Bus) Unsubscribe(id AgentID) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if ch, ok := b.channels[id]; ok {
		close(ch)
		delete(b.channels, id)
		b.logger.Info("bus_unsubscribe", "agent", id)
	}
}