ai-operator/internal/asterisk/ari/websocket.go

package ari

import (
	"context"
	"errors"
	"log/slog"
	"net/http"
	"time"

	"github.com/gorilla/websocket"

	"ai-operator/internal/config"
)

type EventHandler interface {
	HandleEvent(ctx context.Context, event Event) error
}

type EventListener struct {
	cfg                       config.AsteriskConfig
	mode                      WSAuthMode
	handler                   EventHandler
	logger                    *slog.Logger
	dialer                    *websocket.Dialer
	stopOnApplicationReplaced bool
}

func NewEventListener(cfg config.AsteriskConfig, mode WSAuthMode, handler EventHandler, logger *slog.Logger) *EventListener {
	if mode == "" {
		mode = WSAuthMode(cfg.ARIWSAuthMode)
	}
	if mode == "" {
		mode = WSAuthBasic
	}
	return &EventListener{cfg: cfg, mode: mode, handler: handler, logger: logger, dialer: &websocket.Dialer{HandshakeTimeout: 5 * time.Second}, stopOnApplicationReplaced: true}
}

func (l *EventListener) Run(ctx context.Context) error {
	backoffs := []time.Duration{time.Second, 2 * time.Second, 5 * time.Second, 10 * time.Second}
	attempt := 0
	for {
		if err := ctx.Err(); err != nil {
			return nil
		}
		err := l.runOnce(ctx)
		if err == nil || errors.Is(err, context.Canceled) {
			return nil
		}
		if errors.Is(err, ErrApplicationReplaced) {
			return err
		}
		d := backoffs[min(attempt, len(backoffs)-1)]
		attempt++
		l.logWarn("ari websocket disconnected", "error", err, "reconnect_in", d.String())
		select {
		case <-ctx.Done():
			return nil
		case <-time.After(d):
		}
	}
}

var ErrApplicationReplaced = errors.New("ARI application replaced by another websocket")

func (l *EventListener) runOnce(ctx context.Context) error {
	wsURL, err := BuildWebSocketURL(l.cfg, l.mode)
	if err != nil {
		return err
	}
	header := http.Header{}
	if wsURL.AuthMode == WSAuthBasic || wsURL.AuthMode == WSAuthAuto {
		req, _ := http.NewRequest(http.MethodGet, wsURL.URL, nil)
		req.SetBasicAuth(l.cfg.ARIUser, l.cfg.ARIPassword)
		header = req.Header
	}
	l.logInfo("connecting ari websocket", "url", wsURL.Sanitized, "auth_mode", string(wsURL.AuthMode))
	conn, resp, err := l.dialer.DialContext(ctx, wsURL.URL, header)
	if err != nil && wsURL.AuthMode == WSAuthAuto && resp != nil && resp.StatusCode == http.StatusUnauthorized {
		fallback, ferr := BuildWebSocketURL(l.cfg, WSAuthQueryAPIKey)
		if ferr != nil {
			return ferr
		}
		l.logWarn("basic auth rejected, trying query api_key fallback", "url", fallback.Sanitized)
		conn, _, err = l.dialer.DialContext(ctx, fallback.URL, nil)
	}
	if err != nil {
		return err
	}
	defer conn.Close()
	l.logInfo("ari websocket connected", "app", l.cfg.ARIApp)
	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			return err
		}
		event, err := ParseEvent(data)
		if err != nil {
			l.logWarn("invalid ari event json", "error", err)
			continue
		}
		if event.EventType() == EventApplicationReplaced {
			l.logWarn("ari application replaced", "app", l.cfg.ARIApp)
			if l.stopOnApplicationReplaced {
				return ErrApplicationReplaced
			}
		}
		if l.handler != nil {
			if err := l.handler.HandleEvent(ctx, event); err != nil {
				l.logWarn("ari event handler error", "event_type", string(event.EventType()), "error", err)
			}
		}
	}
}

func (l *EventListener) logInfo(msg string, args ...any) {
	if l.logger != nil {
		l.logger.Info(msg, args...)
	}
}
func (l *EventListener) logWarn(msg string, args ...any) {
	if l.logger != nil {
		l.logger.Warn(msg, args...)
	}
}
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func CheckWebSocketConnect(ctx context.Context, cfg config.AsteriskConfig, mode WSAuthMode) error {
	wsURL, err := BuildWebSocketURL(cfg, mode)
	if err != nil {
		return err
	}
	header := http.Header{}
	if wsURL.AuthMode == WSAuthBasic || wsURL.AuthMode == WSAuthAuto {
		req, _ := http.NewRequest(http.MethodGet, wsURL.URL, nil)
		req.SetBasicAuth(cfg.ARIUser, cfg.ARIPassword)
		header = req.Header
	}
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL.URL, header)
	if err != nil {
		return err
	}
	return conn.Close()
}