diff --git a/config/config.example.json b/config/config.example.json index 4153d44bc..b4ccfb781 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -68,6 +68,13 @@ "reconnect_interval": 5, "group_trigger_prefix": [], "allow_from": [] + }, + "websocket": { + "enabled": false, + "host": "127.0.0.1", + "port": 18793, + "path": "/ws", + "allow_from": [] } }, "providers": { diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index 7f6abc4cb..11a6fc2c0 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -176,6 +176,19 @@ func (m *Manager) initChannels() error { } } + if m.config.Channels.WebSocket.Enabled { + logger.DebugC("channels", "Attempting to initialize WebSocket channel") + ws, err := NewWebSocketChannel(m.config.Channels.WebSocket, m.bus) + if err != nil { + logger.ErrorCF("channels", "Failed to initialize WebSocket channel", map[string]interface{}{ + "error": err.Error(), + }) + } else { + m.channels["websocket"] = ws + logger.InfoC("channels", "WebSocket channel enabled successfully") + } + } + logger.InfoCF("channels", "Channel initialization completed", map[string]interface{}{ "enabled_channels": len(m.channels), }) diff --git a/pkg/channels/websocket.go b/pkg/channels/websocket.go new file mode 100644 index 000000000..dab08a64f --- /dev/null +++ b/pkg/channels/websocket.go @@ -0,0 +1,289 @@ +package channels + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "os" + "path/filepath" + "sync" + + "github.com/google/uuid" + "github.com/gorilla/websocket" + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" +) + +// wsIncoming is the JSON message sent from APK to picoclaw. +type wsIncoming struct { + Content string `json:"content"` + SenderID string `json:"sender_id,omitempty"` + Images []string `json:"images,omitempty"` +} + +// wsOutgoing is the JSON message sent from picoclaw to APK. +type wsOutgoing struct { + Content string `json:"content"` +} + +// WebSocketChannel is a server-side WebSocket channel that accepts +// connections from clients (e.g. a Google Assistant replacement APK). +type WebSocketChannel struct { + *BaseChannel + config config.WebSocketConfig + server *http.Server + upgrader websocket.Upgrader + clients map[*websocket.Conn]string // conn → clientID + chatConns map[string]*websocket.Conn // chatID → conn + mu sync.RWMutex + ctx context.Context + cancel context.CancelFunc +} + +func NewWebSocketChannel(cfg config.WebSocketConfig, msgBus *bus.MessageBus) (*WebSocketChannel, error) { + base := NewBaseChannel("websocket", cfg, msgBus, cfg.AllowFrom) + + return &WebSocketChannel{ + BaseChannel: base, + config: cfg, + upgrader: websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + }, + clients: make(map[*websocket.Conn]string), + chatConns: make(map[string]*websocket.Conn), + }, nil +} + +func (c *WebSocketChannel) Start(ctx context.Context) error { + logger.InfoC("websocket", "Starting WebSocket channel server") + + c.ctx, c.cancel = context.WithCancel(ctx) + + mux := http.NewServeMux() + mux.HandleFunc(c.config.Path, c.handleWS) + + addr := fmt.Sprintf("%s:%d", c.config.Host, c.config.Port) + c.server = &http.Server{ + Addr: addr, + Handler: mux, + } + + c.setRunning(true) + + logger.InfoCF("websocket", "WebSocket server listening", map[string]interface{}{ + "host": c.config.Host, + "port": c.config.Port, + "path": c.config.Path, + }) + + go func() { + if err := c.server.ListenAndServe(); err != nil && err != http.ErrServerClosed { + logger.ErrorCF("websocket", "Server error", map[string]interface{}{ + "error": err.Error(), + }) + } + }() + + return nil +} + +func (c *WebSocketChannel) Stop(ctx context.Context) error { + logger.InfoC("websocket", "Stopping WebSocket channel") + c.setRunning(false) + + if c.cancel != nil { + c.cancel() + } + + c.mu.Lock() + for conn, clientID := range c.clients { + logger.DebugCF("websocket", "Closing client connection", map[string]interface{}{ + "client_id": clientID, + }) + conn.Close() + } + c.clients = make(map[*websocket.Conn]string) + c.chatConns = make(map[string]*websocket.Conn) + c.mu.Unlock() + + if c.server != nil { + if err := c.server.Shutdown(ctx); err != nil { + logger.ErrorCF("websocket", "Server shutdown error", map[string]interface{}{ + "error": err.Error(), + }) + return err + } + } + + logger.InfoC("websocket", "WebSocket channel stopped") + return nil +} + +func (c *WebSocketChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { + if !c.IsRunning() { + return fmt.Errorf("websocket channel not running") + } + + c.mu.RLock() + conn, ok := c.chatConns[msg.ChatID] + c.mu.RUnlock() + + if !ok { + return fmt.Errorf("no connection for chat %s", msg.ChatID) + } + + out := wsOutgoing{Content: msg.Content} + data, err := json.Marshal(out) + if err != nil { + return fmt.Errorf("failed to marshal response: %w", err) + } + + c.mu.RLock() + defer c.mu.RUnlock() + + // Verify connection still exists (may have been removed during cleanup). + if _, exists := c.clients[conn]; !exists { + return fmt.Errorf("connection for chat %s no longer active", msg.ChatID) + } + + if err := conn.WriteMessage(websocket.TextMessage, data); err != nil { + logger.ErrorCF("websocket", "Failed to send message", map[string]interface{}{ + "chat_id": msg.ChatID, + "error": err.Error(), + }) + return err + } + + return nil +} + +func (c *WebSocketChannel) handleWS(w http.ResponseWriter, r *http.Request) { + conn, err := c.upgrader.Upgrade(w, r, nil) + if err != nil { + logger.ErrorCF("websocket", "Upgrade failed", map[string]interface{}{ + "error": err.Error(), + }) + return + } + + clientID := uuid.New().String() + + logger.InfoCF("websocket", "New WebSocket connection", map[string]interface{}{ + "client_id": clientID, + "remote_addr": r.RemoteAddr, + }) + + chatID := fmt.Sprintf("ws:%s", clientID) + + c.mu.Lock() + c.clients[conn] = clientID + c.chatConns[chatID] = conn + c.mu.Unlock() + + go c.readPump(conn, clientID, chatID) +} + +func (c *WebSocketChannel) readPump(conn *websocket.Conn, clientID, chatID string) { + defer func() { + c.mu.Lock() + delete(c.clients, conn) + delete(c.chatConns, chatID) + c.mu.Unlock() + conn.Close() + + logger.InfoCF("websocket", "Client disconnected", map[string]interface{}{ + "client_id": clientID, + }) + }() + + for { + select { + case <-c.ctx.Done(): + return + default: + } + + _, message, err := conn.ReadMessage() + if err != nil { + if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { + logger.ErrorCF("websocket", "Read error", map[string]interface{}{ + "client_id": clientID, + "error": err.Error(), + }) + } + return + } + + var incoming wsIncoming + if err := json.Unmarshal(message, &incoming); err != nil { + logger.ErrorCF("websocket", "Invalid JSON message", map[string]interface{}{ + "client_id": clientID, + "error": err.Error(), + }) + continue + } + + // Use sender_id from message if provided, otherwise use clientID. + senderID := clientID + if incoming.SenderID != "" { + senderID = incoming.SenderID + } + + content := incoming.Content + var media []string + + // Save images to temp files (same pattern as Telegram). + for i, imgData := range incoming.Images { + path, err := c.saveImage(imgData) + if err != nil { + logger.ErrorCF("websocket", "Failed to save image", map[string]interface{}{ + "client_id": clientID, + "index": i, + "error": err.Error(), + }) + continue + } + media = append(media, path) + content += fmt.Sprintf("\n[image: photo_%d]", i) + } + + logger.DebugCF("websocket", "Received message", map[string]interface{}{ + "client_id": clientID, + "content": incoming.Content, + "images": len(incoming.Images), + }) + + c.HandleMessage(senderID, chatID, content, media, nil) + } +} + +func (c *WebSocketChannel) saveImage(base64Data string) (string, error) { + data, err := base64.StdEncoding.DecodeString(base64Data) + if err != nil { + // Try URL-safe base64. + data, err = base64.URLEncoding.DecodeString(base64Data) + if err != nil { + return "", fmt.Errorf("failed to decode base64: %w", err) + } + } + + tmpDir := filepath.Join(os.TempDir(), "picoclaw", "ws_images") + if err := os.MkdirAll(tmpDir, 0755); err != nil { + return "", fmt.Errorf("failed to create temp dir: %w", err) + } + + f, err := os.CreateTemp(tmpDir, "ws_img_*.png") + if err != nil { + return "", fmt.Errorf("failed to create temp file: %w", err) + } + defer f.Close() + + if _, err := f.Write(data); err != nil { + return "", fmt.Errorf("failed to write image: %w", err) + } + + return f.Name(), nil +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 2c6378555..259293902 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -71,16 +71,17 @@ type AgentDefaults struct { } type ChannelsConfig struct { - WhatsApp WhatsAppConfig `json:"whatsapp"` - Telegram TelegramConfig `json:"telegram"` - Feishu FeishuConfig `json:"feishu"` - Discord DiscordConfig `json:"discord"` - MaixCam MaixCamConfig `json:"maixcam"` - QQ QQConfig `json:"qq"` - DingTalk DingTalkConfig `json:"dingtalk"` - Slack SlackConfig `json:"slack"` - LINE LINEConfig `json:"line"` - OneBot OneBotConfig `json:"onebot"` + WhatsApp WhatsAppConfig `json:"whatsapp"` + Telegram TelegramConfig `json:"telegram"` + Feishu FeishuConfig `json:"feishu"` + Discord DiscordConfig `json:"discord"` + MaixCam MaixCamConfig `json:"maixcam"` + QQ QQConfig `json:"qq"` + DingTalk DingTalkConfig `json:"dingtalk"` + Slack SlackConfig `json:"slack"` + LINE LINEConfig `json:"line"` + OneBot OneBotConfig `json:"onebot"` + WebSocket WebSocketConfig `json:"websocket"` } type WhatsAppConfig struct { @@ -158,6 +159,14 @@ type OneBotConfig struct { AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_ONEBOT_ALLOW_FROM"` } +type WebSocketConfig struct { + Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WEBSOCKET_ENABLED"` + Host string `json:"host" env:"PICOCLAW_CHANNELS_WEBSOCKET_HOST"` + Port int `json:"port" env:"PICOCLAW_CHANNELS_WEBSOCKET_PORT"` + Path string `json:"path" env:"PICOCLAW_CHANNELS_WEBSOCKET_PATH"` + AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WEBSOCKET_ALLOW_FROM"` +} + type HeartbeatConfig struct { Enabled bool `json:"enabled" env:"PICOCLAW_HEARTBEAT_ENABLED"` Interval int `json:"interval" env:"PICOCLAW_HEARTBEAT_INTERVAL"` // minutes, min 5 @@ -316,6 +325,13 @@ func DefaultConfig() *Config { GroupTriggerPrefix: []string{}, AllowFrom: FlexibleStringSlice{}, }, + WebSocket: WebSocketConfig{ + Enabled: true, + Host: "127.0.0.1", + Port: 18793, + Path: "/ws", + AllowFrom: FlexibleStringSlice{}, + }, }, Providers: ProvidersConfig{ Anthropic: ProviderConfig{},