feat: implement WebSocket channel (TDD green phase)
Replace stubs with full implementation: connect/auth handshake, listen loop with bus publishing, send with content-type detection, reconnect loop, and clean shutdown. All 14 tests pass with -race. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
d49723e380
commit
a92d421c5b
1 changed files with 182 additions and 1 deletions
|
|
@ -3,13 +3,16 @@ package websocket
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"sync"
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
"github.com/gorilla/websocket"
|
"github.com/gorilla/websocket"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
"github.com/sipeed/picoclaw/pkg/channels"
|
"github.com/sipeed/picoclaw/pkg/channels"
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
)
|
)
|
||||||
|
|
||||||
// WSEnvelope is the typed envelope for all WebSocket messages.
|
// WSEnvelope is the typed envelope for all WebSocket messages.
|
||||||
|
|
@ -46,6 +49,7 @@ type WSOutboundMessage struct {
|
||||||
type WebSocketChannel struct {
|
type WebSocketChannel struct {
|
||||||
*channels.BaseChannel
|
*channels.BaseChannel
|
||||||
config config.WebSocketConfig
|
config config.WebSocketConfig
|
||||||
|
mb *bus.MessageBus
|
||||||
conn *websocket.Conn
|
conn *websocket.Conn
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
cancel context.CancelFunc
|
cancel context.CancelFunc
|
||||||
|
|
@ -53,7 +57,7 @@ type WebSocketChannel struct {
|
||||||
writeMu sync.Mutex // guards writes
|
writeMu sync.Mutex // guards writes
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewWebSocketChannel creates a new WebSocket channel (stub).
|
// NewWebSocketChannel creates a new WebSocket channel.
|
||||||
func NewWebSocketChannel(cfg config.WebSocketConfig, mb *bus.MessageBus) (*WebSocketChannel, error) {
|
func NewWebSocketChannel(cfg config.WebSocketConfig, mb *bus.MessageBus) (*WebSocketChannel, error) {
|
||||||
base := channels.NewBaseChannel(
|
base := channels.NewBaseChannel(
|
||||||
"websocket",
|
"websocket",
|
||||||
|
|
@ -65,15 +69,130 @@ func NewWebSocketChannel(cfg config.WebSocketConfig, mb *bus.MessageBus) (*WebSo
|
||||||
ch := &WebSocketChannel{
|
ch := &WebSocketChannel{
|
||||||
BaseChannel: base,
|
BaseChannel: base,
|
||||||
config: cfg,
|
config: cfg,
|
||||||
|
mb: mb,
|
||||||
}
|
}
|
||||||
return ch, nil
|
return ch, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// connect dials the WebSocket server and sends the auth handshake.
|
||||||
|
func (c *WebSocketChannel) connect() error {
|
||||||
|
conn, _, err := websocket.DefaultDialer.Dial(c.config.WSUrl, nil)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("websocket dial: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Send auth envelope.
|
||||||
|
authData, err := json.Marshal(WSAuthData{
|
||||||
|
AgentID: c.config.AgentID,
|
||||||
|
Token: c.config.AccessToken,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
conn.Close()
|
||||||
|
return fmt.Errorf("marshal auth: %w", err)
|
||||||
|
}
|
||||||
|
env := WSEnvelope{Type: "auth", Data: json.RawMessage(authData)}
|
||||||
|
envBytes, err := json.Marshal(env)
|
||||||
|
if err != nil {
|
||||||
|
conn.Close()
|
||||||
|
return fmt.Errorf("marshal auth envelope: %w", err)
|
||||||
|
}
|
||||||
|
if err := conn.WriteMessage(websocket.TextMessage, envBytes); err != nil {
|
||||||
|
conn.Close()
|
||||||
|
return fmt.Errorf("write auth: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
c.conn = conn
|
||||||
|
c.mu.Unlock()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// listen reads messages from the WebSocket connection and publishes them to the bus.
|
||||||
|
func (c *WebSocketChannel) listen() {
|
||||||
|
c.mu.Lock()
|
||||||
|
conn := c.conn
|
||||||
|
c.mu.Unlock()
|
||||||
|
if conn == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
for {
|
||||||
|
_, data, err := conn.ReadMessage()
|
||||||
|
if err != nil {
|
||||||
|
c.mu.Lock()
|
||||||
|
c.conn = nil
|
||||||
|
c.mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
env, err := parseEnvelope(data)
|
||||||
|
if err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if env.Type != "message" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var msg WSInboundMessage
|
||||||
|
if err := json.Unmarshal(env.Data, &msg); err != nil {
|
||||||
|
logger.ErrorCF("websocket", "Failed to unmarshal inbound message", map[string]any{
|
||||||
|
"error": err.Error(),
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
inbound := bus.InboundMessage{
|
||||||
|
Channel: "websocket",
|
||||||
|
SenderID: msg.SenderID,
|
||||||
|
ChatID: msg.ChatID,
|
||||||
|
Content: msg.Content,
|
||||||
|
MessageID: msg.MessageID,
|
||||||
|
SessionKey: msg.SessionKey,
|
||||||
|
Metadata: msg.Metadata,
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := c.mb.PublishInbound(c.ctx, inbound); err != nil {
|
||||||
|
logger.ErrorCF("websocket", "Failed to publish inbound", map[string]any{
|
||||||
|
"error": err.Error(),
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (c *WebSocketChannel) Start(ctx context.Context) error {
|
func (c *WebSocketChannel) Start(ctx context.Context) error {
|
||||||
|
c.ctx, c.cancel = context.WithCancel(ctx)
|
||||||
|
|
||||||
|
if err := c.connect(); err != nil {
|
||||||
|
if c.config.ReconnectInterval <= 0 {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
logger.ErrorCF("websocket", "Initial connect failed, will retry", map[string]any{
|
||||||
|
"error": err.Error(),
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
go c.listen()
|
||||||
|
}
|
||||||
|
|
||||||
|
if c.config.ReconnectInterval > 0 {
|
||||||
|
go c.reconnectLoop()
|
||||||
|
}
|
||||||
|
|
||||||
|
c.SetRunning(true)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *WebSocketChannel) Stop(ctx context.Context) error {
|
func (c *WebSocketChannel) Stop(ctx context.Context) error {
|
||||||
|
c.SetRunning(false)
|
||||||
|
c.cancel()
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
if c.conn != nil {
|
||||||
|
c.conn.Close()
|
||||||
|
c.conn = nil
|
||||||
|
}
|
||||||
|
c.mu.Unlock()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -81,9 +200,71 @@ func (c *WebSocketChannel) Send(ctx context.Context, msg bus.OutboundMessage) er
|
||||||
if !c.IsRunning() {
|
if !c.IsRunning() {
|
||||||
return channels.ErrNotRunning
|
return channels.ErrNotRunning
|
||||||
}
|
}
|
||||||
|
|
||||||
|
out := WSOutboundMessage{
|
||||||
|
Channel: msg.Channel,
|
||||||
|
ChatID: msg.ChatID,
|
||||||
|
Content: msg.Content,
|
||||||
|
ContentType: DetectContentType(msg.Content),
|
||||||
|
}
|
||||||
|
outData, err := json.Marshal(out)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("marshal outbound: %w", err)
|
||||||
|
}
|
||||||
|
env := WSEnvelope{Type: "message", Data: json.RawMessage(outData)}
|
||||||
|
envBytes, err := json.Marshal(env)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("marshal envelope: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
conn := c.conn
|
||||||
|
c.mu.Unlock()
|
||||||
|
if conn == nil {
|
||||||
|
return channels.ErrTemporary
|
||||||
|
}
|
||||||
|
|
||||||
|
c.writeMu.Lock()
|
||||||
|
err = conn.WriteMessage(websocket.TextMessage, envBytes)
|
||||||
|
c.writeMu.Unlock()
|
||||||
|
if err != nil {
|
||||||
|
return channels.ErrTemporary
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// reconnectLoop periodically reconnects when the connection is lost.
|
||||||
|
func (c *WebSocketChannel) reconnectLoop() {
|
||||||
|
interval := time.Duration(c.config.ReconnectInterval) * time.Second
|
||||||
|
if interval < time.Second {
|
||||||
|
interval = time.Second
|
||||||
|
}
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-c.ctx.Done():
|
||||||
|
return
|
||||||
|
case <-time.After(interval):
|
||||||
|
}
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
needsReconnect := c.conn == nil
|
||||||
|
c.mu.Unlock()
|
||||||
|
|
||||||
|
if !needsReconnect {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := c.connect(); err != nil {
|
||||||
|
logger.ErrorCF("websocket", "Reconnect failed", map[string]any{
|
||||||
|
"error": err.Error(),
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
go c.listen()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// DetectContentType returns "json" if content is valid JSON, otherwise "text".
|
// DetectContentType returns "json" if content is valid JSON, otherwise "text".
|
||||||
func DetectContentType(content string) string {
|
func DetectContentType(content string) string {
|
||||||
if json.Valid([]byte(content)) {
|
if json.Valid([]byte(content)) {
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue