feat: add WebSocket server channel for APK integration
Add a WebSocket channel that accepts connections from clients (e.g. Google Assistant replacement APK). Text + base64 images are received via JSON, images are saved locally following the existing media pattern. Enabled by default on 127.0.0.1:18793/ws. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
KoheiYamashita
parent
add20a4454
commit
a639d0ecdc
4 changed files with 335 additions and 10 deletions
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
})
|
||||
|
|
|
|||
289
pkg/channels/websocket.go
Normal file
289
pkg/channels/websocket.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -81,6 +81,7 @@ type ChannelsConfig struct {
|
|||
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{},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue