first commit

transports/bifrost-http/websocket/connection.go

@@ -0,0 +1,140 @@
+// Package websocket provides upstream WebSocket connection management for the Bifrost gateway.
+// It manages pooled connections to provider WebSocket APIs (e.g., OpenAI Responses WS mode,
+// Realtime API) and client session bindings.
+package websocket
+
+import (
+	"net/http"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	ws "github.com/fasthttp/websocket"
+	"github.com/maximhq/bifrost/core/schemas"
+)
+
+// UpstreamConn wraps a WebSocket connection to an upstream provider.
+// Thread-safe for concurrent read/write via separate mutexes.
+type UpstreamConn struct {
+	conn      *ws.Conn
+	provider  schemas.ModelProvider
+	keyID     string
+	endpoint  string
+	createdAt time.Time
+	lastUsed  atomic.Int64 // unix nano
+
+	writeMu sync.Mutex
+	readMu  sync.Mutex
+
+	closed atomic.Bool
+}
+
+// newUpstreamConn creates a new UpstreamConn wrapping the given websocket connection.
+func newUpstreamConn(conn *ws.Conn, provider schemas.ModelProvider, keyID, endpoint string) *UpstreamConn {
+	uc := &UpstreamConn{
+		conn:      conn,
+		provider:  provider,
+		keyID:     keyID,
+		endpoint:  endpoint,
+		createdAt: time.Now(),
+	}
+	uc.lastUsed.Store(time.Now().UnixNano())
+	return uc
+}
+
+// WriteMessage sends a message to the upstream provider. Thread-safe.
+func (c *UpstreamConn) WriteMessage(messageType int, data []byte) error {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	c.lastUsed.Store(time.Now().UnixNano())
+	return c.conn.WriteMessage(messageType, data)
+}
+
+// WriteJSON sends a JSON-encoded message to the upstream provider. Thread-safe.
+func (c *UpstreamConn) WriteJSON(v interface{}) error {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	c.lastUsed.Store(time.Now().UnixNano())
+	return c.conn.WriteJSON(v)
+}
+
+// ReadMessage reads a message from the upstream provider. Thread-safe.
+func (c *UpstreamConn) ReadMessage() (messageType int, p []byte, err error) {
+	c.readMu.Lock()
+	defer c.readMu.Unlock()
+	c.lastUsed.Store(time.Now().UnixNano())
+	return c.conn.ReadMessage()
+}
+
+// Close closes the underlying WebSocket connection.
+func (c *UpstreamConn) Close() error {
+	if c.closed.CompareAndSwap(false, true) {
+		return c.conn.Close()
+	}
+	return nil
+}
+
+// IsClosed returns whether the connection has been closed.
+func (c *UpstreamConn) IsClosed() bool {
+	return c.closed.Load()
+}
+
+// Provider returns the provider this connection is for.
+func (c *UpstreamConn) Provider() schemas.ModelProvider {
+	return c.provider
+}
+
+// KeyID returns the API key ID used for this connection.
+func (c *UpstreamConn) KeyID() string {
+	return c.keyID
+}
+
+// CreatedAt returns when this connection was established.
+func (c *UpstreamConn) CreatedAt() time.Time {
+	return c.createdAt
+}
+
+// LastUsed returns the last time this connection was used.
+func (c *UpstreamConn) LastUsed() time.Time {
+	return time.Unix(0, c.lastUsed.Load())
+}
+
+// Age returns how long this connection has been alive.
+func (c *UpstreamConn) Age() time.Duration {
+	return time.Since(c.createdAt)
+}
+
+// SetReadDeadline sets the read deadline on the underlying connection.
+func (c *UpstreamConn) SetReadDeadline(t time.Time) error {
+	return c.conn.SetReadDeadline(t)
+}
+
+// SetWriteDeadline sets the write deadline on the underlying connection.
+func (c *UpstreamConn) SetWriteDeadline(t time.Time) error {
+	return c.conn.SetWriteDeadline(t)
+}
+
+// SetPongHandler sets a handler for pong messages received from the upstream.
+func (c *UpstreamConn) SetPongHandler(h func(appData string) error) {
+	c.conn.SetPongHandler(h)
+}
+
+// WritePing sends a ping message to the upstream. Thread-safe.
+func (c *UpstreamConn) WritePing(data []byte) error {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	c.lastUsed.Store(time.Now().UnixNano())
+	return c.conn.WriteMessage(ws.PingMessage, data)
+}
+
+// Dial creates a new WebSocket connection to the given URL with the provided headers.
+func Dial(url string, headers map[string]string) (*ws.Conn, *http.Response, error) {
+	dialer := ws.Dialer{
+		HandshakeTimeout: 10 * time.Second,
+	}
+	h := http.Header{}
+	for k, v := range headers {
+		h.Set(k, v)
+	}
+	return dialer.Dial(url, h)
+}