diff --git a/pkg/channels/websocket/channel.go b/pkg/channels/websocket/channel.go new file mode 100644 index 000000000..e06bea39f --- /dev/null +++ b/pkg/channels/websocket/channel.go @@ -0,0 +1,417 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// WebSocket channel implementation for local LAN chat + +package websocket + +import ( + "context" + _ "embed" + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/gorilla/websocket" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" +) + +//go:embed chat.html +var chatHTML []byte + +//go:embed logo.jpg +var logoImage []byte + +// Channel implements the Channel interface for WebSocket connections +type Channel struct { + config config.WebSocketConfig + bus *bus.MessageBus + running bool + allowList []string + server *http.Server + upgrader websocket.Upgrader + clients sync.Map // map[string]*websocket.Conn + clientsMu sync.RWMutex + ctx context.Context + cancel context.CancelFunc +} + +// WebSocketMessage represents the JSON message format +type WebSocketMessage struct { + Type string `json:"type"` // "chat", "status", "error", "system" + Content string `json:"content"` // Message content + Sender string `json:"sender"` // Sender ID + Timestamp int64 `json:"timestamp"` // Unix timestamp in milliseconds + SessionID string `json:"session_id,omitempty"` +} + +// NewChannel creates a new WebSocket channel instance +func NewChannel(cfg config.WebSocketConfig, messageBus *bus.MessageBus) (*Channel, error) { + if cfg.Port == 0 { + cfg.Port = 8080 + } + if cfg.Host == "" { + cfg.Host = "0.0.0.0" + } + + return &Channel{ + config: cfg, + bus: messageBus, + allowList: cfg.AllowFrom, + running: false, + upgrader: websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + // Allow all origins for local LAN usage + // For production, you should implement proper origin checking + return true + }, + ReadBufferSize: 1024, + WriteBufferSize: 1024, + }, + }, nil +} + +// Name returns the channel name +func (c *Channel) Name() string { + return "websocket" +} + +// IsRunning returns whether the channel is currently running +func (c *Channel) IsRunning() bool { + c.clientsMu.RLock() + defer c.clientsMu.RUnlock() + return c.running +} + +// IsAllowed checks if a sender ID is allowed to use this channel +func (c *Channel) IsAllowed(senderID string) bool { + if len(c.allowList) == 0 { + return true + } + for _, allowed := range c.allowList { + if strings.EqualFold(allowed, senderID) { + return true + } + } + return false +} + +// setRunning sets the running state +func (c *Channel) setRunning(running bool) { + c.clientsMu.Lock() + defer c.clientsMu.Unlock() + c.running = running +} + +// HandleMessage processes an incoming message and publishes it to the bus +func (c *Channel) HandleMessage(senderID, chatID, content string, media []string, metadata map[string]string) { + if !c.IsAllowed(senderID) { + return + } + + // Build session key: channel:chatID + sessionKey := fmt.Sprintf("%s:%s", c.Name(), chatID) + + msg := bus.InboundMessage{ + Channel: c.Name(), + SenderID: senderID, + ChatID: chatID, + Content: content, + Media: media, + SessionKey: sessionKey, + Metadata: metadata, + } + + c.bus.PublishInbound(msg) +} + +// Start initializes and starts the WebSocket server +func (c *Channel) Start(ctx context.Context) error { + c.ctx, c.cancel = context.WithCancel(ctx) + + mux := http.NewServeMux() + mux.HandleFunc("/ws", c.handleWebSocket) + mux.HandleFunc("/", c.handleIndex) + mux.HandleFunc("/assets/logo.jpg", c.handleLogo) + + addr := fmt.Sprintf("%s:%d", c.config.Host, c.config.Port) + c.server = &http.Server{ + Addr: addr, + Handler: mux, + ReadTimeout: 15 * time.Second, + WriteTimeout: 15 * time.Second, + } + + c.setRunning(true) + logger.InfoCF("websocket", "WebSocket channel starting", map[string]any{ + "address": addr, + }) + + // Start server in goroutine + errCh := make(chan error, 1) + go func() { + if err := c.server.ListenAndServe(); err != nil && err != http.ErrServerClosed { + errCh <- err + } + }() + + // Check for immediate startup errors + select { + case err := <-errCh: + c.setRunning(false) + return fmt.Errorf("failed to start WebSocket server: %w", err) + case <-time.After(100 * time.Millisecond): + logger.InfoCF("websocket", "WebSocket channel started successfully", map[string]any{ + "address": addr, + }) + return nil + } +} + +// Stop gracefully shuts down the WebSocket server +func (c *Channel) Stop(ctx context.Context) error { + logger.InfoC("websocket", "Stopping WebSocket channel") + + if c.cancel != nil { + c.cancel() + } + + // Close all client connections + c.clients.Range(func(key, value any) bool { + if conn, ok := value.(*websocket.Conn); ok { + conn.Close() + } + c.clients.Delete(key) + return true + }) + + // Shutdown HTTP server + if c.server != nil { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := c.server.Shutdown(shutdownCtx); err != nil { + logger.ErrorCF("websocket", "Error shutting down server", map[string]any{ + "error": err.Error(), + }) + } + } + + c.setRunning(false) + logger.InfoC("websocket", "WebSocket channel stopped") + return nil +} + +// Send sends a message to the specified chat (client) +func (c *Channel) Send(ctx context.Context, msg bus.OutboundMessage) error { + if !c.IsRunning() { + return fmt.Errorf("websocket channel not running") + } + + wsMsg := WebSocketMessage{ + Type: "chat", + Content: msg.Content, + Sender: "assistant", + Timestamp: time.Now().UnixMilli(), + } + + data, err := json.Marshal(wsMsg) + if err != nil { + return fmt.Errorf("failed to marshal message: %w", err) + } + + // Ensure UTF-8 validity + if !json.Valid(data) { + logger.ErrorCF("websocket", "Invalid JSON data", map[string]any{ + "content_preview": msg.Content[:min(len(msg.Content), 100)], + }) + return fmt.Errorf("invalid JSON message") + } + + // Send to specific client or broadcast + if msg.ChatID != "" && msg.ChatID != "broadcast" { + // Send to specific client + if conn, ok := c.clients.Load(msg.ChatID); ok { + if wsConn, ok := conn.(*websocket.Conn); ok { + err := wsConn.WriteMessage(websocket.TextMessage, data) + if err != nil { + // Connection may be dead, clean it up + logger.WarnCF("websocket", "Failed to send to client, removing connection", map[string]any{ + "client": msg.ChatID, + "error": err.Error(), + }) + c.clients.Delete(msg.ChatID) + return err + } + return nil + } + } + return fmt.Errorf("client %s not found", msg.ChatID) + } + + // Broadcast to all connected clients + var lastErr error + deadClients := make([]any, 0) + c.clients.Range(func(key, value any) bool { + if conn, ok := value.(*websocket.Conn); ok { + if err := conn.WriteMessage(websocket.TextMessage, data); err != nil { + logger.WarnCF("websocket", "Failed to send to client", map[string]any{ + "client": key, + "error": err.Error(), + }) + deadClients = append(deadClients, key) + lastErr = err + } + } + return true + }) + + // Clean up dead connections + for _, key := range deadClients { + c.clients.Delete(key) + logger.InfoCF("websocket", "Removed dead client connection", map[string]any{ + "client": key, + }) + } + + return lastErr +} + +// Helper function for min +func min(a, b int) int { + if a < b { + return a + } + return b +} + +// handleWebSocket handles WebSocket connection upgrades +func (c *Channel) handleWebSocket(w http.ResponseWriter, r *http.Request) { + conn, err := c.upgrader.Upgrade(w, r, nil) + if err != nil { + logger.ErrorCF("websocket", "Failed to upgrade connection", map[string]any{ + "error": err.Error(), + }) + return + } + + // Generate client ID from remote address + clientID := r.RemoteAddr + c.clients.Store(clientID, conn) + + logger.InfoCF("websocket", "New client connected", map[string]any{ + "client_id": clientID, + }) + + // Handle client messages + go c.handleClient(clientID, conn) +} + +// handleClient processes messages from a WebSocket client +func (c *Channel) handleClient(clientID string, conn *websocket.Conn) { + defer func() { + conn.Close() + c.clients.Delete(clientID) + logger.InfoCF("websocket", "Client disconnected", map[string]any{ + "client_id": clientID, + }) + }() + + // Set up ping/pong handlers for connection health check + conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + conn.SetPongHandler(func(string) error { + conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + return nil + }) + + // Start ping ticker + pingTicker := time.NewTicker(30 * time.Second) + defer pingTicker.Stop() + + // Channel for messages + messageChan := make(chan []byte, 10) + defer close(messageChan) + + // Read messages in a goroutine + go func() { + for { + _, message, err := conn.ReadMessage() + if err != nil { + if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { + logger.ErrorCF("websocket", "WebSocket error", map[string]any{ + "client_id": clientID, + "error": err.Error(), + }) + } + return + } + messageChan <- message + } + }() + + for { + select { + case <-c.ctx.Done(): + return + case <-pingTicker.C: + // Send ping to check connection health + if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(10*time.Second)); err != nil { + logger.WarnCF("websocket", "Failed to send ping", map[string]any{ + "client_id": clientID, + "error": err.Error(), + }) + return + } + case message, ok := <-messageChan: + if !ok { + return + } + + var wsMsg WebSocketMessage + if err := json.Unmarshal(message, &wsMsg); err != nil { + logger.WarnCF("websocket", "Failed to parse message", map[string]any{ + "client_id": clientID, + "error": err.Error(), + }) + continue + } + + // Process chat messages + if wsMsg.Type == "chat" { + // Check allowlist + if !c.IsAllowed(clientID) { + logger.WarnCF("websocket", "Unauthorized client", map[string]any{ + "client_id": clientID, + }) + continue + } + + logger.DebugCF("websocket", "Received message", map[string]any{ + "client_id": clientID, + "content": wsMsg.Content, + }) + + // Send to agent via message bus + c.HandleMessage(clientID, clientID, wsMsg.Content, nil, nil) + } + } + } +} + +// handleLogo serves the logo image +func (c *Channel) handleLogo(w http.ResponseWriter, r *http.Request) { + // Use embedded logo image + w.Header().Set("Content-Type", "image/jpeg") + w.Header().Set("Cache-Control", "public, max-age=86400") // Cache for 1 day + w.Write(logoImage) +} + +// handleIndex serves a simple HTML chat interface +func (c *Channel) handleIndex(w http.ResponseWriter, r *http.Request) { + // Use embedded HTML file + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.Write(chatHTML) +} diff --git a/pkg/channels/websocket/chat.html b/pkg/channels/websocket/chat.html new file mode 100644 index 000000000..0becc0154 --- /dev/null +++ b/pkg/channels/websocket/chat.html @@ -0,0 +1,1267 @@ + + + + + + + + + PicoClaw · Intelligent Conversation + + + + + + + +
+
+
+ +

PicoClaw WebSocket Chat

+
+
+
+
+ Connecting... +
+
+ EN +
+
+ +
+
+
+ +
+
+
+
+
+
+
+ +
+
+ +
+
+ + + +
+ +
+
+
+
+ + + + + + diff --git a/pkg/channels/websocket/logo.jpg b/pkg/channels/websocket/logo.jpg new file mode 100644 index 000000000..7f1a7c2b7 Binary files /dev/null and b/pkg/channels/websocket/logo.jpg differ diff --git a/pkg/channels/websocket_adapter.go b/pkg/channels/websocket_adapter.go new file mode 100644 index 000000000..2cec1719a --- /dev/null +++ b/pkg/channels/websocket_adapter.go @@ -0,0 +1,18 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package channels + +import ( + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/channels/websocket" + "github.com/sipeed/picoclaw/pkg/config" +) + +// NewWebSocketChannel creates a new WebSocket channel instance. +// This is an adapter to maintain naming consistency with other channel constructors. +func NewWebSocketChannel(cfg config.WebSocketConfig, messageBus *bus.MessageBus) (Channel, error) { + return websocket.NewChannel(cfg, messageBus) +}