diff --git a/docs/zh/chat-apps.md b/docs/zh/chat-apps.md index aeba7d460..e9c102870 100644 --- a/docs/zh/chat-apps.md +++ b/docs/zh/chat-apps.md @@ -26,7 +26,7 @@ PicoClaw 支持多种聊天平台，使您的 Agent 能够连接到任何地方 | **IRC** | ⭐⭐ 中等 | 服务器 + TLS 配置 | [查看文档](#irc) | | **OneBot** | ⭐⭐ 中等 | 兼容 NapCat/Go-CQHTTP，社区生态丰富 | [查看文档](../channels/onebot/README.zh.md) | | **MaixCam** | ⭐ 简单 | 专为 AI 摄像头设计的硬件集成通道 | [查看文档](../channels/maixcam/README.zh.md) | -| **Pico** | ⭐ 简单 | PicoClaw 原生协议通道 | | +| **Pico** | ⭐ 简单 | PicoClaw 原生协议通道（内置 Web 聊天） | [本节](#pico) | --- @@ -663,3 +663,19 @@ picoclaw gateway ``` + + + +
+Pico（内置 Web 聊天） + +Pico 是 PicoClaw 原生协议通道，用于自带 Web UI 与 Agent 对话。 + +* **主路径**：浏览器通过 **WebSocket** 连接 `GET /pico/ws`（由 Web 服务反代到 Gateway，与 `ws_url` 一致）。 +* **降级**：若 WebSocket 无法建立（超时或失败），前端会自动改用 **SSE** 接收推送：`GET /pico/events?session_id=...`，并使用 **`POST /pico/send`** 发送用户消息。两者均使用请求头 `Authorization: Bearer `，token 与 `GET /api/pico/token` 返回的相同；同一次响应中还会提供 `events_url`、`send_url` 供客户端使用。 +* **反向代理**：若 Web UI 前有 **nginx** 等代理，请对 `/pico/events` **关闭响应缓冲**（例如 `proxy_buffering off;`），否则 SSE 可能被缓冲导致无法实时显示流式回复。 +* **单会话连接**：同一 `session_id` 在服务端只保留 **一条** 实时下行连接（新的 WebSocket 或 SSE 会断开同会话的旧连接），避免重复推送；前端对相同 `message_id` 的 `message.create` 也会合并为一条展示。 + +启用方式见项目配置中的 `channels.pico`；运行 `picoclaw gateway` 并打开 Web 控制台即可使用内置聊天。 + +
diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index 86ce98b06..66e75da8a 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -40,6 +40,20 @@ func (pc *picoConn) writeJSON(v any) error { return pc.conn.WriteJSON(v) } +// ID returns the connection id (picoSubscriber). +func (pc *picoConn) ID() string { return pc.id } + +// SessionID returns the bound session (picoSubscriber). +func (pc *picoConn) SessionID() string { return pc.sessionID } + +// Deliver sends a Pico message over the WebSocket (picoSubscriber). +func (pc *picoConn) Deliver(msg PicoMessage) error { + return pc.writeJSON(msg) +} + +// Close shuts down the WebSocket (picoSubscriber). +func (pc *picoConn) Close() { pc.close() } + // close closes the connection. func (pc *picoConn) close() { if pc.closed.CompareAndSwap(false, true) { @@ -50,13 +64,13 @@ func (pc *picoConn) close() { } } -// PicoChannel implements the native Pico Protocol WebSocket channel. +// PicoChannel implements the native Pico Protocol channel (WebSocket plus HTTP SSE fallback). // It serves as the reference implementation for all optional capability interfaces. type PicoChannel struct { *channels.BaseChannel config config.PicoConfig upgrader websocket.Upgrader - connections sync.Map // connID → *picoConn + subscribers sync.Map // connID → picoSubscriber (*picoConn or *picoSSEConn) connCount atomic.Int32 ctx context.Context cancel context.CancelFunc @@ -109,12 +123,12 @@ func (c *PicoChannel) Stop(ctx context.Context) error { logger.InfoC("pico", "Stopping Pico Protocol channel") c.SetRunning(false) - // Close all connections - c.connections.Range(func(key, value any) bool { - if pc, ok := value.(*picoConn); ok { - pc.close() + // Close all subscribers (WebSocket + SSE) + c.subscribers.Range(func(key, value any) bool { + if sub, ok := value.(picoSubscriber); ok { + sub.Close() } - c.connections.Delete(key) + c.subscribers.Delete(key) return true }) @@ -136,6 +150,18 @@ func (c *PicoChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) { switch { case path == "/ws" || path == "/ws/": c.handleWebSocket(w, r) + case path == "/events" || path == "/events/": + if r.Method == http.MethodGet { + c.handleSSE(w, r) + } else { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + } + case path == "/send" || path == "/send/": + if r.Method == http.MethodPost { + c.handlePostSend(w, r) + } else { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + } default: http.NotFound(w, r) } @@ -208,20 +234,21 @@ func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error { msg.SessionID = sessionID var sent bool - c.connections.Range(func(key, value any) bool { - pc, ok := value.(*picoConn) + c.subscribers.Range(func(key, value any) bool { + sub, ok := value.(picoSubscriber) if !ok { return true } - if pc.sessionID == sessionID { - if err := pc.writeJSON(msg); err != nil { - logger.DebugCF("pico", "Write to connection failed", map[string]any{ - "conn_id": pc.id, - "error": err.Error(), - }) - } else { - sent = true - } + if sub.SessionID() != sessionID { + return true + } + if err := sub.Deliver(msg); err != nil { + logger.DebugCF("pico", "Write to subscriber failed", map[string]any{ + "conn_id": sub.ID(), + "error": err.Error(), + }) + } else { + sent = true } return true }) @@ -232,6 +259,25 @@ func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error { return nil } +// disconnectSubscribersForSession closes and removes every subscriber for the given session so +// a new WebSocket or SSE connection does not overlap with the previous one (avoids duplicate pushes). +func (c *PicoChannel) disconnectSubscribersForSession(sessionID string) { + c.subscribers.Range(func(key, value any) bool { + sub, ok := value.(picoSubscriber) + if !ok || sub.SessionID() != sessionID { + return true + } + c.subscribers.Delete(key) + c.connCount.Add(-1) + sub.Close() + logger.InfoCF("pico", "Disconnected prior subscriber (single session)", map[string]any{ + "conn_id": sub.ID(), + "session_id": sessionID, + }) + return true + }) +} + // handleWebSocket upgrades the HTTP connection and manages the WebSocket lifecycle. func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) { if !c.IsRunning() { @@ -245,6 +291,13 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) { return } + sessionID := r.URL.Query().Get("session_id") + if sessionID == "" { + sessionID = uuid.New().String() + } + // Drop any prior connection for this session before enforcing the global cap. + c.disconnectSubscribersForSession(sessionID) + // Check connection limit maxConns := c.config.MaxConnections if maxConns <= 0 { @@ -269,19 +322,13 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) { return } - // Determine session ID from query param or generate one - sessionID := r.URL.Query().Get("session_id") - if sessionID == "" { - sessionID = uuid.New().String() - } - pc := &picoConn{ id: uuid.New().String(), conn: conn, sessionID: sessionID, } - c.connections.Store(pc.id, pc) + c.subscribers.Store(pc.id, pc) c.connCount.Add(1) logger.InfoCF("pico", "WebSocket client connected", map[string]any{ @@ -341,8 +388,9 @@ func (c *PicoChannel) matchedSubprotocol(r *http.Request) string { func (c *PicoChannel) readLoop(pc *picoConn) { defer func() { pc.close() - c.connections.Delete(pc.id) - c.connCount.Add(-1) + if _, loaded := c.subscribers.LoadAndDelete(pc.id); loaded { + c.connCount.Add(-1) + } logger.InfoCF("pico", "WebSocket client disconnected", map[string]any{ "conn_id": pc.id, "session_id": pc.sessionID, @@ -452,6 +500,11 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) { sessionID = pc.sessionID } + c.publishUserMessage(pc.id, sessionID, msg.ID, content) +} + +// publishUserMessage forwards validated user text into the message bus. +func (c *PicoChannel) publishUserMessage(connID, sessionID, msgID, content string) { chatID := "pico:" + sessionID senderID := "pico-user" @@ -460,7 +513,7 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) { metadata := map[string]string{ "platform": "pico", "session_id": sessionID, - "conn_id": pc.id, + "conn_id": connID, } logger.DebugCF("pico", "Received message", map[string]any{ @@ -478,7 +531,7 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) { return } - c.HandleMessage(c.ctx, peer, msg.ID, senderID, chatID, content, nil, metadata, sender) + c.HandleMessage(c.ctx, peer, msgID, senderID, chatID, content, nil, metadata, sender) } // truncate truncates a string to maxLen runes. diff --git a/pkg/channels/pico/pico_http_test.go b/pkg/channels/pico/pico_http_test.go new file mode 100644 index 000000000..c7699db9e --- /dev/null +++ b/pkg/channels/pico/pico_http_test.go @@ -0,0 +1,224 @@ +package pico + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" +) + +func TestPicoChannel_PostSend_Unauthorized(t *testing.T) { + t.Parallel() + ch, err := NewPicoChannel(config.PicoConfig{Token: "tok"}, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := ch.Start(ctx); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ch.Stop(ctx) }) + + body := `{"type":"message.send","session_id":"s1","id":"1","payload":{"content":"hi"}}` + req := httptest.NewRequest(http.MethodPost, "/pico/send", strings.NewReader(body)) + rec := httptest.NewRecorder() + ch.ServeHTTP(rec, req) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("code = %d, want 401", rec.Code) + } +} + +func TestPicoChannel_PostSend_EmptyContent(t *testing.T) { + t.Parallel() + ch, err := NewPicoChannel(config.PicoConfig{Token: "tok"}, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := ch.Start(ctx); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ch.Stop(ctx) }) + + body := `{"type":"message.send","session_id":"s1","id":"1","payload":{"content":" "}}` + req := httptest.NewRequest(http.MethodPost, "/pico/send", strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer tok") + rec := httptest.NewRecorder() + ch.ServeHTTP(rec, req) + if rec.Code != http.StatusBadRequest { + t.Fatalf("code = %d, want 400", rec.Code) + } + var m map[string]string + if err := json.Unmarshal(rec.Body.Bytes(), &m); err != nil { + t.Fatal(err) + } + if m["error"] != "empty_content" { + t.Fatalf("error = %q, want empty_content", m["error"]) + } +} + +func TestPicoChannel_PostSend_MissingSession(t *testing.T) { + t.Parallel() + ch, err := NewPicoChannel(config.PicoConfig{Token: "tok"}, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := ch.Start(ctx); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ch.Stop(ctx) }) + + body := `{"type":"message.send","id":"1","payload":{"content":"hi"}}` + req := httptest.NewRequest(http.MethodPost, "/pico/send", strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer tok") + rec := httptest.NewRecorder() + ch.ServeHTTP(rec, req) + if rec.Code != http.StatusBadRequest { + t.Fatalf("code = %d, want 400", rec.Code) + } +} + +func TestPicoChannel_PostSend_OK(t *testing.T) { + t.Parallel() + ch, err := NewPicoChannel(config.PicoConfig{Token: "tok"}, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := ch.Start(ctx); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ch.Stop(ctx) }) + + body := `{"type":"message.send","session_id":"s1","id":"m1","payload":{"content":"hello"}}` + req := httptest.NewRequest(http.MethodPost, "/pico/send", strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer tok") + rec := httptest.NewRecorder() + ch.ServeHTTP(rec, req) + if rec.Code != http.StatusNoContent { + t.Fatalf("code = %d, want 204", rec.Code) + } +} + +func TestPicoChannel_GetEvents_Unauthorized(t *testing.T) { + t.Parallel() + ch, err := NewPicoChannel(config.PicoConfig{Token: "tok"}, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := ch.Start(ctx); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ch.Stop(ctx) }) + + req := httptest.NewRequest(http.MethodGet, "/pico/events?session_id=s1", nil) + rec := httptest.NewRecorder() + ch.ServeHTTP(rec, req) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("code = %d, want 401", rec.Code) + } +} + +func TestPicoChannel_SSE_ReceivesBroadcast(t *testing.T) { + ch, err := NewPicoChannel(config.PicoConfig{Token: "tok"}, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := ch.Start(ctx); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ch.Stop(ctx) }) + + ctxReq, cancel := context.WithCancel(context.Background()) + req := httptest.NewRequest(http.MethodGet, "/pico/events?session_id=sse-s1", nil) + req = req.WithContext(ctxReq) + req.Header.Set("Authorization", "Bearer tok") + rec := httptest.NewRecorder() + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + ch.ServeHTTP(rec, req) + }() + + time.Sleep(200 * time.Millisecond) + + if err := ch.Send(ctx, bus.OutboundMessage{ChatID: "pico:sse-s1", Content: "broadcast-body"}); err != nil { + t.Fatalf("Send: %v", err) + } + + time.Sleep(150 * time.Millisecond) + cancel() + wg.Wait() + + out := rec.Body.String() + if !strings.Contains(out, "broadcast-body") { + t.Fatalf("expected outbound in body, got: %q", out) + } + if !strings.Contains(out, "message.create") { + t.Fatalf("expected message type in stream, got: %q", out) + } +} + +func TestPicoChannel_SSE_SecondConnectionSameSessionReplacesFirst(t *testing.T) { + ch, err := NewPicoChannel(config.PicoConfig{Token: "tok"}, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := ch.Start(ctx); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ch.Stop(ctx) }) + + const sess = "shared-sse" + + ctx1, cancel1 := context.WithCancel(context.Background()) + req1 := httptest.NewRequest(http.MethodGet, "/pico/events?session_id="+sess, nil).WithContext(ctx1) + req1.Header.Set("Authorization", "Bearer tok") + rec1 := httptest.NewRecorder() + go ch.ServeHTTP(rec1, req1) + time.Sleep(200 * time.Millisecond) + + ctx2, cancel2 := context.WithCancel(context.Background()) + req2 := httptest.NewRequest(http.MethodGet, "/pico/events?session_id="+sess, nil).WithContext(ctx2) + req2.Header.Set("Authorization", "Bearer tok") + rec2 := httptest.NewRecorder() + var wg2 sync.WaitGroup + wg2.Add(1) + go func() { + defer wg2.Done() + ch.ServeHTTP(rec2, req2) + }() + time.Sleep(200 * time.Millisecond) + + if err := ch.Send(ctx, bus.OutboundMessage{ChatID: "pico:" + sess, Content: "only-once"}); err != nil { + t.Fatalf("Send: %v", err) + } + time.Sleep(150 * time.Millisecond) + + n1 := strings.Count(rec1.Body.String(), "only-once") + n2 := strings.Count(rec2.Body.String(), "only-once") + if n2 != 1 { + t.Fatalf("second SSE should receive exactly one broadcast, got n2=%d", n2) + } + if n1 != 0 { + t.Fatalf("first SSE should be replaced and not receive broadcast, got n1=%d", n1) + } + + cancel2() + wg2.Wait() + cancel1() + time.Sleep(100 * time.Millisecond) +} diff --git a/pkg/channels/pico/pico_sse.go b/pkg/channels/pico/pico_sse.go new file mode 100644 index 000000000..a0fc518cc --- /dev/null +++ b/pkg/channels/pico/pico_sse.go @@ -0,0 +1,221 @@ +package pico + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/google/uuid" + + "github.com/sipeed/picoclaw/pkg/logger" +) + +// picoSubscriber is a WebSocket or SSE client registered for outbound Pico messages. +type picoSubscriber interface { + ID() string + SessionID() string + Deliver(msg PicoMessage) error + Close() +} + +// picoSSEConn streams outbound messages as Server-Sent Events. +type picoSSEConn struct { + id string + sessionID string + w http.ResponseWriter + rc *http.ResponseController + writeMu sync.Mutex + closed atomic.Bool + cancel context.CancelFunc +} + +func (s *picoSSEConn) ID() string { return s.id } + +func (s *picoSSEConn) SessionID() string { return s.sessionID } + +func (s *picoSSEConn) Deliver(msg PicoMessage) error { + if s.closed.Load() { + return fmt.Errorf("connection closed") + } + b, err := json.Marshal(msg) + if err != nil { + return err + } + s.writeMu.Lock() + defer s.writeMu.Unlock() + if s.closed.Load() { + return fmt.Errorf("connection closed") + } + if _, err := fmt.Fprintf(s.w, "data: %s\n\n", b); err != nil { + return err + } + return s.rc.Flush() +} + +func (s *picoSSEConn) Close() { + s.shutdown() +} + +func (s *picoSSEConn) shutdown() { + if s.closed.CompareAndSwap(false, true) { + if s.cancel != nil { + s.cancel() + } + } +} + +func (c *PicoChannel) handleSSE(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if !c.IsRunning() { + http.Error(w, "channel not running", http.StatusServiceUnavailable) + return + } + if !c.authenticate(r) { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + + sessionID := r.URL.Query().Get("session_id") + if sessionID == "" { + sessionID = uuid.New().String() + } + c.disconnectSubscribersForSession(sessionID) + + maxConns := c.config.MaxConnections + if maxConns <= 0 { + maxConns = 100 + } + if int(c.connCount.Load()) >= maxConns { + http.Error(w, "too many connections", http.StatusServiceUnavailable) + return + } + + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + + rc := http.NewResponseController(w) + + id := uuid.New().String() + ctx, cancel := context.WithCancel(r.Context()) + sse := &picoSSEConn{ + id: id, + sessionID: sessionID, + w: w, + rc: rc, + cancel: cancel, + } + + c.subscribers.Store(id, sse) + c.connCount.Add(1) + defer func() { + sse.shutdown() + if _, loaded := c.subscribers.LoadAndDelete(id); loaded { + c.connCount.Add(-1) + } + logger.InfoCF("pico", "SSE client disconnected", map[string]any{ + "conn_id": id, + "session_id": sessionID, + }) + }() + + logger.InfoCF("pico", "SSE client connected", map[string]any{ + "conn_id": id, + "session_id": sessionID, + }) + + ready, err := json.Marshal(map[string]string{"conn_id": id, "session_id": sessionID}) + if err != nil { + return + } + sse.writeMu.Lock() + _, werr := fmt.Fprintf(w, "event: ready\ndata: %s\n\n", ready) + sse.writeMu.Unlock() + if werr != nil { + return + } + if err := rc.Flush(); err != nil { + return + } + + pingInterval := time.Duration(c.config.PingInterval) * time.Second + if pingInterval <= 0 { + pingInterval = 30 * time.Second + } + ticker := time.NewTicker(pingInterval) + defer ticker.Stop() + + for { + select { + case <-c.ctx.Done(): + return + case <-ctx.Done(): + return + case <-ticker.C: + sse.writeMu.Lock() + if sse.closed.Load() { + sse.writeMu.Unlock() + return + } + _, werr := fmt.Fprintf(w, ": ping\n\n") + sse.writeMu.Unlock() + if werr != nil { + return + } + if err := rc.Flush(); err != nil { + return + } + } + } +} + +func (c *PicoChannel) handlePostSend(w http.ResponseWriter, r *http.Request) { + if !c.IsRunning() { + http.Error(w, "channel not running", http.StatusServiceUnavailable) + return + } + if !c.authenticate(r) { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + + var msg PicoMessage + if err := json.NewDecoder(r.Body).Decode(&msg); err != nil { + writePicoSendError(w, http.StatusBadRequest, "invalid_json") + return + } + if msg.Type != TypeMessageSend { + writePicoSendError(w, http.StatusBadRequest, "unsupported_type") + return + } + content, _ := msg.Payload["content"].(string) + if strings.TrimSpace(content) == "" { + writePicoSendError(w, http.StatusBadRequest, "empty_content") + return + } + sessionID := strings.TrimSpace(msg.SessionID) + if sessionID == "" { + writePicoSendError(w, http.StatusBadRequest, "session_id_required") + return + } + + connID := "http-" + uuid.New().String() + c.publishUserMessage(connID, sessionID, msg.ID, content) + w.WriteHeader(http.StatusNoContent) +} + +func writePicoSendError(w http.ResponseWriter, code int, errCode string) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + _ = json.NewEncoder(w).Encode(map[string]string{"error": errCode}) +} diff --git a/web/backend/api/pico.go b/web/backend/api/pico.go index 8fbb8737f..b4798d280 100644 --- a/web/backend/api/pico.go +++ b/web/backend/api/pico.go @@ -22,6 +22,9 @@ func (h *Handler) registerPicoRoutes(mux *http.ServeMux) { // This allows the frontend to connect via the same port as the web UI, // avoiding the need to expose extra ports for WebSocket communication. mux.HandleFunc("GET /pico/ws", h.handleWebSocketProxy()) + // SSE and HTTP send fallback (same gateway Pico channel). + mux.HandleFunc("GET /pico/events", h.handleGatewayProxy()) + mux.HandleFunc("POST /pico/send", h.handleGatewayProxy()) } // createWsProxy creates a reverse proxy to the current gateway WebSocket endpoint. @@ -37,6 +40,11 @@ func (h *Handler) createWsProxy() *httputil.ReverseProxy { // handleWebSocketProxy wraps a reverse proxy to handle WebSocket connections. // The reverse proxy forwards the incoming upgrade handshake as-is. func (h *Handler) handleWebSocketProxy() http.HandlerFunc { + return h.handleGatewayProxy() +} + +// handleGatewayProxy forwards /pico/ws, /pico/events, and /pico/send to the gateway HTTP server. +func (h *Handler) handleGatewayProxy() http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { proxy := h.createWsProxy() proxy.ServeHTTP(w, r) @@ -54,12 +62,16 @@ func (h *Handler) handleGetPicoToken(w http.ResponseWriter, r *http.Request) { } wsURL := h.buildWsURL(r, cfg) + eventsURL := h.buildPicoEventsURL(r, cfg) + sendURL := h.buildPicoSendURL(r, cfg) w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]any{ - "token": cfg.Channels.Pico.Token(), - "ws_url": wsURL, - "enabled": cfg.Channels.Pico.Enabled, + "token": cfg.Channels.Pico.Token, + "ws_url": wsURL, + "events_url": eventsURL, + "send_url": sendURL, + "enabled": cfg.Channels.Pico.Enabled, }) } @@ -82,11 +94,15 @@ func (h *Handler) handleRegenPicoToken(w http.ResponseWriter, r *http.Request) { } wsURL := h.buildWsURL(r, cfg) + eventsURL := h.buildPicoEventsURL(r, cfg) + sendURL := h.buildPicoSendURL(r, cfg) w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]any{ - "token": token, - "ws_url": wsURL, + "token": token, + "ws_url": wsURL, + "events_url": eventsURL, + "send_url": sendURL, }) } @@ -147,13 +163,17 @@ func (h *Handler) handlePicoSetup(w http.ResponseWriter, r *http.Request) { } wsURL := h.buildWsURL(r, cfg) + eventsURL := h.buildPicoEventsURL(r, cfg) + sendURL := h.buildPicoSendURL(r, cfg) w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]any{ - "token": cfg.Channels.Pico.Token(), - "ws_url": wsURL, - "enabled": true, - "changed": changed, + "token": cfg.Channels.Pico.Token, + "ws_url": wsURL, + "events_url": eventsURL, + "send_url": sendURL, + "enabled": true, + "changed": changed, }) } diff --git a/web/frontend/src/api/pico.ts b/web/frontend/src/api/pico.ts index 9a1a553d5..da15d72e9 100644 --- a/web/frontend/src/api/pico.ts +++ b/web/frontend/src/api/pico.ts @@ -3,12 +3,16 @@ interface PicoTokenResponse { token: string ws_url: string + events_url?: string + send_url?: string enabled: boolean } interface PicoSetupResponse { token: string ws_url: string + events_url?: string + send_url?: string enabled: boolean changed: boolean } diff --git a/web/frontend/src/components/chat/chat-page.tsx b/web/frontend/src/components/chat/chat-page.tsx index ebcde8981..91c6e91a8 100644 --- a/web/frontend/src/components/chat/chat-page.tsx +++ b/web/frontend/src/components/chat/chat-page.tsx @@ -79,9 +79,9 @@ export function ChatPage() { } }, [messages, isTyping, isAtBottom]) - const handleSend = () => { + const handleSend = async () => { if (!input.trim() || !canSend) return - if (sendMessage(input.trim())) { + if (await sendMessage(input.trim())) { setInput("") } } diff --git a/web/frontend/src/features/chat/controller.ts b/web/frontend/src/features/chat/controller.ts index 5e6eb2229..dfd141045 100644 --- a/web/frontend/src/features/chat/controller.ts +++ b/web/frontend/src/features/chat/controller.ts @@ -1,12 +1,14 @@ import { getDefaultStore } from "jotai" import { toast } from "sonner" +import { launcherFetch } from "@/api/http" import { getPicoToken } from "@/api/pico" import { loadSessionMessages, mergeHistoryMessages, } from "@/features/chat/history" import { type PicoMessage, handlePicoMessage } from "@/features/chat/protocol" +import { parsePicoSSEData, readPicoSSEStream } from "@/features/chat/sse" import { clearStoredSessionId, generateSessionId, @@ -23,6 +25,8 @@ import { type GatewayState, gatewayAtom } from "@/store/gateway" const store = getDefaultStore() +const WS_CONNECT_TIMEOUT_MS = 8000 + let wsRef: WebSocket | null = null let isConnecting = false let msgIdCounter = 0 @@ -35,6 +39,12 @@ let reconnectTimer: number | null = null let reconnectAttempts = 0 let shouldMaintainConnection = false +/** SSE long-poll stream active (fetch returned 200 and body is being read). */ +let sseActive = false +let sseAbort: AbortController | null = null +let picoTokenRef: string | null = null +let picoSendUrlRef: string | null = null + function clearReconnectTimer() { if (reconnectTimer !== null) { window.clearTimeout(reconnectTimer) @@ -42,6 +52,14 @@ function clearReconnectTimer() { } } +function stopSSE() { + sseAbort?.abort() + sseAbort = null + sseActive = false + picoSendUrlRef = null + picoTokenRef = null +} + function shouldReconnectFor(generation: number, sessionId: string): boolean { return ( shouldMaintainConnection && @@ -73,8 +91,8 @@ function needsActiveSessionHydration(): boolean { return Boolean( storedSessionId && - storedSessionId === state.activeSessionId && - !state.hasHydratedActiveSession, + storedSessionId === state.activeSessionId && + !state.hasHydratedActiveSession, ) } @@ -100,6 +118,7 @@ function disconnectChatInternal({ isConnecting = false invalidateSocket(socket) + stopSSE() updateChatStore({ connectionState: "disconnected", @@ -107,6 +126,228 @@ function disconnectChatInternal({ }) } +function attachWebSocketHandlers( + socket: WebSocket, + generation: number, + sessionId: string, +) { + socket.onmessage = (event) => { + if ( + !isCurrentSocket({ + socket, + currentSocket: wsRef, + generation, + currentGeneration: connectionGeneration, + sessionId, + currentSessionId: activeSessionIdRef, + }) + ) { + return + } + + try { + const message = JSON.parse(event.data) as PicoMessage + handlePicoMessage(message, sessionId) + } catch { + console.warn("Non-JSON message from pico:", event.data) + } + } + + socket.onclose = () => { + if ( + !isCurrentSocket({ + socket, + currentSocket: wsRef, + generation, + currentGeneration: connectionGeneration, + sessionId, + currentSessionId: activeSessionIdRef, + }) + ) { + return + } + wsRef = null + isConnecting = false + updateChatStore({ + connectionState: "disconnected", + isTyping: false, + }) + scheduleReconnect(generation, sessionId) + } + + socket.onerror = () => { + if ( + !isCurrentSocket({ + socket, + currentSocket: wsRef, + generation, + currentGeneration: connectionGeneration, + sessionId, + currentSessionId: activeSessionIdRef, + }) + ) { + return + } + isConnecting = false + updateChatStore({ connectionState: "error" }) + scheduleReconnect(generation, sessionId) + } +} + +/** + * Try WebSocket first. Resolves false if connection fails or times out (then caller may use SSE). + */ +function tryOpenWebSocket( + generation: number, + sessionId: string, + token: string, + wsUrl: string, +): Promise { + return new Promise((resolve) => { + const finalWsUrl = normalizeWsUrlForBrowser(wsUrl) + const url = `${finalWsUrl}?session_id=${encodeURIComponent(sessionId)}` + const socket = new WebSocket(url, [`token.${token}`]) + + if (generation !== connectionGeneration) { + invalidateSocket(socket) + resolve(false) + return + } + + let settled = false + const finishFail = () => { + if (settled) { + return + } + settled = true + window.clearTimeout(timer) + invalidateSocket(socket) + try { + socket.close() + } catch { + /* ignore */ + } + resolve(false) + } + + const timer = window.setTimeout(() => { + if (socket.readyState !== WebSocket.OPEN) { + finishFail() + } + }, WS_CONNECT_TIMEOUT_MS) + + socket.onerror = () => { + if (socket.readyState !== WebSocket.OPEN) { + finishFail() + } + } + + socket.onopen = () => { + window.clearTimeout(timer) + if (settled || generation !== connectionGeneration) { + invalidateSocket(socket) + resolve(false) + return + } + settled = true + wsRef = socket + updateChatStore({ connectionState: "connected" }) + isConnecting = false + reconnectAttempts = 0 + attachWebSocketHandlers(socket, generation, sessionId) + resolve(true) + } + }) +} + +async function openSSETransport( + generation: number, + sessionId: string, + token: string, + eventsUrlRaw: string | undefined, + sendUrlRaw: string | undefined, +): Promise { + const eventsPath = eventsUrlRaw?.trim() || "/pico/events" + const sendPath = sendUrlRaw?.trim() || "/pico/send" + + const eventsURL = new URL(eventsPath, window.location.origin) + eventsURL.searchParams.set("session_id", sessionId) + + const sendURL = new URL(sendPath, window.location.origin) + + picoTokenRef = token + picoSendUrlRef = sendURL.toString() + + sseAbort = new AbortController() + const signal = sseAbort.signal + + let res: Response + try { + res = await launcherFetch(eventsURL.toString(), { + headers: { Authorization: `Bearer ${token}` }, + signal, + credentials: "same-origin", + }) + } catch { + if (generation !== connectionGeneration) { + return + } + stopSSE() + isConnecting = false + updateChatStore({ connectionState: "error" }) + scheduleReconnect(generation, sessionId) + return + } + + if (generation !== connectionGeneration) { + stopSSE() + return + } + + if (!res.ok || !res.body) { + stopSSE() + isConnecting = false + updateChatStore({ connectionState: "error" }) + scheduleReconnect(generation, sessionId) + return + } + + sseActive = true + isConnecting = false + reconnectAttempts = 0 + updateChatStore({ connectionState: "connected" }) + + try { + await readPicoSSEStream(res.body, signal, (data) => { + if (generation !== connectionGeneration) { + return + } + const message = parsePicoSSEData(data) + if (message) { + handlePicoMessage(message, sessionId) + } + }) + } catch { + /* aborted or read error */ + } finally { + sseActive = false + sseAbort = null + picoSendUrlRef = null + picoTokenRef = null + + if ( + generation === connectionGeneration && + shouldReconnectFor(generation, sessionId) + ) { + updateChatStore({ + connectionState: "disconnected", + isTyping: false, + }) + scheduleReconnect(generation, sessionId) + } + } +} + export async function connectChat() { if ( store.get(gatewayAtom).status !== "running" || @@ -115,11 +356,16 @@ export async function connectChat() { return } + if (isConnecting) { + return + } + if (sseActive) { + return + } if ( - isConnecting || - (wsRef && - (wsRef.readyState === WebSocket.OPEN || - wsRef.readyState === WebSocket.CONNECTING)) + wsRef && + (wsRef.readyState === WebSocket.OPEN || + wsRef.readyState === WebSocket.CONNECTING) ) { return } @@ -128,11 +374,17 @@ export async function connectChat() { connectionGeneration = generation isConnecting = true clearReconnectTimer() + + invalidateSocket(wsRef) + wsRef = null + stopSSE() + updateChatStore({ connectionState: "connecting" }) + const sessionId = activeSessionIdRef + try { - const { token, ws_url } = await getPicoToken() - const sessionId = activeSessionIdRef + const { token, ws_url, events_url, send_url } = await getPicoToken() if (generation !== connectionGeneration) { isConnecting = false @@ -147,103 +399,32 @@ export async function connectChat() { return } - const finalWsUrl = normalizeWsUrlForBrowser(ws_url) - const url = `${finalWsUrl}?session_id=${encodeURIComponent(sessionId)}` - const socket = new WebSocket(url, [`token.${token}`]) + const wsOk = await tryOpenWebSocket( + generation, + sessionId, + token, + ws_url, + ) if (generation !== connectionGeneration) { isConnecting = false - invalidateSocket(socket) return } - socket.onopen = () => { - if ( - !isCurrentSocket({ - socket, - currentSocket: wsRef, - generation, - currentGeneration: connectionGeneration, - sessionId, - currentSessionId: activeSessionIdRef, - }) - ) { - return - } - updateChatStore({ connectionState: "connected" }) - isConnecting = false - reconnectAttempts = 0 + if (wsOk) { + return } - socket.onmessage = (event) => { - if ( - !isCurrentSocket({ - socket, - currentSocket: wsRef, - generation, - currentGeneration: connectionGeneration, - sessionId, - currentSessionId: activeSessionIdRef, - }) - ) { - return - } - - try { - const message = JSON.parse(event.data) as PicoMessage - handlePicoMessage(message, sessionId) - } catch { - console.warn("Non-JSON message from pico:", event.data) - } - } - - socket.onclose = () => { - if ( - !isCurrentSocket({ - socket, - currentSocket: wsRef, - generation, - currentGeneration: connectionGeneration, - sessionId, - currentSessionId: activeSessionIdRef, - }) - ) { - return - } - wsRef = null - isConnecting = false - updateChatStore({ - connectionState: "disconnected", - isTyping: false, - }) - scheduleReconnect(generation, sessionId) - } - - socket.onerror = () => { - if ( - !isCurrentSocket({ - socket, - currentSocket: wsRef, - generation, - currentGeneration: connectionGeneration, - sessionId, - currentSessionId: activeSessionIdRef, - }) - ) { - return - } - isConnecting = false - updateChatStore({ connectionState: "error" }) - scheduleReconnect(generation, sessionId) - } - - wsRef = socket + isConnecting = true + updateChatStore({ connectionState: "connecting" }) + await openSSETransport(generation, sessionId, token, events_url, send_url) } catch (error) { if (generation !== connectionGeneration) { isConnecting = false return } console.error("Failed to connect to pico:", error) + stopSSE() updateChatStore({ connectionState: "error" }) isConnecting = false scheduleReconnect(generation, activeSessionIdRef) @@ -324,40 +505,73 @@ export async function hydrateActiveSession() { return hydratePromise } -export function sendChatMessage(content: string) { - if (!wsRef || wsRef.readyState !== WebSocket.OPEN) { - console.warn("WebSocket not connected") - return false - } - - const socket = wsRef +export async function sendChatMessage(content: string): Promise { const id = `msg-${++msgIdCounter}-${Date.now()}` - updateChatStore((prev) => ({ - messages: [ - ...prev.messages, - { id, role: "user", content, timestamp: Date.now() }, - ], - isTyping: true, - })) + const optimistic = () => + updateChatStore((prev) => ({ + messages: [ + ...prev.messages, + { id, role: "user" as const, content, timestamp: Date.now() }, + ], + isTyping: true, + })) - try { - socket.send( - JSON.stringify({ - type: "message.send", - id, - payload: { content }, - }), - ) - return true - } catch (error) { - console.error("Failed to send pico message:", error) + const rollback = () => updateChatStore((prev) => ({ messages: prev.messages.filter((message) => message.id !== id), isTyping: false, })) - return false + + if (wsRef && wsRef.readyState === WebSocket.OPEN) { + optimistic() + try { + wsRef.send( + JSON.stringify({ + type: "message.send", + id, + payload: { content }, + }), + ) + return true + } catch (error) { + console.error("Failed to send pico message:", error) + rollback() + return false + } } + + if (sseActive && picoSendUrlRef && picoTokenRef) { + optimistic() + try { + const res = await launcherFetch(picoSendUrlRef, { + method: "POST", + credentials: "same-origin", + headers: { + Authorization: `Bearer ${picoTokenRef}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + type: "message.send", + id, + session_id: activeSessionIdRef, + payload: { content }, + }), + }) + if (!res.ok) { + rollback() + return false + } + return true + } catch (error) { + console.error("Failed to send pico message:", error) + rollback() + return false + } + } + + console.warn("Pico chat not connected") + return false } export async function switchChatSession(sessionId: string) { diff --git a/web/frontend/src/features/chat/protocol.ts b/web/frontend/src/features/chat/protocol.ts index 5e5220c77..cc5afdcd8 100644 --- a/web/frontend/src/features/chat/protocol.ts +++ b/web/frontend/src/features/chat/protocol.ts @@ -29,18 +29,30 @@ export function handlePicoMessage( ? normalizeUnixTimestamp(Number(message.timestamp)) : Date.now() - updateChatStore((prev) => ({ - messages: [ - ...prev.messages, - { - id: messageId, - role: "assistant", - content, - timestamp, - }, - ], - isTyping: false, - })) + updateChatStore((prev) => { + const exists = prev.messages.some((m) => m.id === messageId) + if (exists) { + return { + ...prev, + messages: prev.messages.map((m) => + m.id === messageId ? { ...m, content, timestamp } : m, + ), + isTyping: false, + } + } + return { + messages: [ + ...prev.messages, + { + id: messageId, + role: "assistant", + content, + timestamp, + }, + ], + isTyping: false, + } + }) break } diff --git a/web/frontend/src/features/chat/sse.ts b/web/frontend/src/features/chat/sse.ts new file mode 100644 index 000000000..e8922b911 --- /dev/null +++ b/web/frontend/src/features/chat/sse.ts @@ -0,0 +1,56 @@ +import type { PicoMessage } from "@/features/chat/protocol" + +/** + * Reads a fetch() response body as text/event-stream and invokes onMessage for each data payload. + * Comment lines (:) are ignored. Supports simple multi-line data: fields (joined with newlines). + */ +export async function readPicoSSEStream( + body: ReadableStream, + signal: AbortSignal, + onMessage: (data: string) => void, +): Promise { + const reader = body.getReader() + const decoder = new TextDecoder() + let buf = "" + + while (!signal.aborted) { + const { done, value } = await reader.read() + if (done) { + break + } + buf += decoder.decode(value, { stream: true }) + + for (;;) { + const sep = buf.indexOf("\n\n") + if (sep < 0) { + break + } + const raw = buf.slice(0, sep) + buf = buf.slice(sep + 2) + + const lines = raw.split("\n") + const dataParts: string[] = [] + for (const line of lines) { + if (line.startsWith(":")) { + continue + } + if (line.startsWith("data:")) { + dataParts.push(line.slice(5).trimStart()) + } + } + const data = dataParts.join("\n").trimEnd() + if (data.length > 0) { + onMessage(data) + } + } + } +} + +export function parsePicoSSEData(data: string): PicoMessage | null { + try { + return JSON.parse(data) as PicoMessage + } catch { + console.warn("Non-JSON SSE data from pico:", data) + return null + } +} diff --git a/web/frontend/vite.config.ts b/web/frontend/vite.config.ts index 0ef4e1415..c1191882d 100644 --- a/web/frontend/vite.config.ts +++ b/web/frontend/vite.config.ts @@ -29,6 +29,11 @@ export default defineConfig({ target: "http://localhost:18800", changeOrigin: true, }, + "/pico": { + target: "http://localhost:18800", + changeOrigin: true, + ws: true, + }, "/ws": { target: "ws://localhost:18800", ws: true,