From c5d5783b292803164fdad7db21b592b196bdf53b Mon Sep 17 00:00:00 2001 From: "zeed.w.zeed" Date: Tue, 24 Mar 2026 09:57:14 +0800 Subject: [PATCH] feat(web-chat): Pico channel SSE fallback and proxy --- docs/zh/chat-apps.md | 18 +- pkg/channels/pico/pico.go | 235 ++++----- pkg/channels/pico/pico_http_test.go | 214 +++++++++ pkg/channels/pico/pico_sse.go | 221 +++++++++ pkg/channels/pico/pico_test.go | 105 +--- web/backend/api/pico.go | 38 +- web/frontend/src/api/pico.ts | 4 + .../src/components/chat/chat-page.tsx | 4 +- web/frontend/src/features/chat/controller.ts | 448 +++++++++++++----- web/frontend/src/features/chat/protocol.ts | 36 +- web/frontend/src/features/chat/sse.ts | 56 +++ web/frontend/vite.config.ts | 5 + 12 files changed, 1008 insertions(+), 376 deletions(-) create mode 100644 pkg/channels/pico/pico_http_test.go create mode 100644 pkg/channels/pico/pico_sse.go create mode 100644 web/frontend/src/features/chat/sse.ts diff --git a/docs/zh/chat-apps.md b/docs/zh/chat-apps.md index 47add38ac..23ea9d4fd 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) | --- @@ -595,3 +595,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 1aa1941cf..548e2d8bc 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,17 +64,16 @@ 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 map[string]*picoConn // connID -> *picoConn - sessionConnections map[string]map[string]*picoConn // sessionID -> connID -> *picoConn - connsMu sync.RWMutex - ctx context.Context - cancel context.CancelFunc + config config.PicoConfig + upgrader websocket.Upgrader + subscribers sync.Map // connID → picoSubscriber (*picoConn or *picoSSEConn) + connCount atomic.Int32 + ctx context.Context + cancel context.CancelFunc } // NewPicoChannel creates a new Pico Protocol channel. @@ -93,104 +106,9 @@ func NewPicoChannel(cfg config.PicoConfig, messageBus *bus.MessageBus) (*PicoCha ReadBufferSize: 1024, WriteBufferSize: 1024, }, - connections: make(map[string]*picoConn), - sessionConnections: make(map[string]map[string]*picoConn), }, nil } -// createAndAddConnection checks MaxConnections and registers a connection atomically. -func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID string, maxConns int) (*picoConn, error) { - c.connsMu.Lock() - defer c.connsMu.Unlock() - if len(c.connections) >= maxConns { - return nil, channels.ErrTemporary - } - - var connID string - for { - connID = uuid.New().String() - if _, exists := c.connections[connID]; !exists { - break - } - } - - pc := &picoConn{ - id: connID, - conn: conn, - sessionID: sessionID, - } - - c.connections[pc.id] = pc - bySession, ok := c.sessionConnections[pc.sessionID] - if !ok { - bySession = make(map[string]*picoConn) - c.sessionConnections[pc.sessionID] = bySession - } - bySession[pc.id] = pc - - return pc, nil -} - -// removeConnection deletes a connection from indexes and returns it when found. -func (c *PicoChannel) removeConnection(connID string) *picoConn { - c.connsMu.Lock() - defer c.connsMu.Unlock() - - pc, ok := c.connections[connID] - if !ok { - return nil - } - - delete(c.connections, connID) - if bySession, ok := c.sessionConnections[pc.sessionID]; ok { - delete(bySession, connID) - if len(bySession) == 0 { - delete(c.sessionConnections, pc.sessionID) - } - } - - return pc -} - -// takeAllConnections snapshots and clears all connection indexes. -func (c *PicoChannel) takeAllConnections() []*picoConn { - c.connsMu.Lock() - defer c.connsMu.Unlock() - - all := make([]*picoConn, 0, len(c.connections)) - for _, pc := range c.connections { - all = append(all, pc) - } - clear(c.connections) - clear(c.sessionConnections) - - return all -} - -// sessionConnectionsSnapshot returns all active connections for a session. -func (c *PicoChannel) sessionConnectionsSnapshot(sessionID string) []*picoConn { - c.connsMu.RLock() - defer c.connsMu.RUnlock() - - bySession, ok := c.sessionConnections[sessionID] - if !ok || len(bySession) == 0 { - return nil - } - - conns := make([]*picoConn, 0, len(bySession)) - for _, pc := range bySession { - conns = append(conns, pc) - } - return conns -} - -// currentConnCount returns a lock-protected snapshot of active connection count. -func (c *PicoChannel) currentConnCount() int { - c.connsMu.RLock() - defer c.connsMu.RUnlock() - return len(c.connections) -} - // Start implements Channel. func (c *PicoChannel) Start(ctx context.Context) error { logger.InfoC("pico", "Starting Pico Protocol channel") @@ -205,10 +123,14 @@ func (c *PicoChannel) Stop(ctx context.Context) error { logger.InfoC("pico", "Stopping Pico Protocol channel") c.SetRunning(false) - // Close all connections - for _, pc := range c.takeAllConnections() { - pc.close() - } + // Close all subscribers (WebSocket + SSE) + c.subscribers.Range(func(key, value any) bool { + if sub, ok := value.(picoSubscriber); ok { + sub.Close() + } + c.subscribers.Delete(key) + return true + }) if c.cancel != nil { c.cancel() @@ -228,6 +150,18 @@ func (c *PicoChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) { switch path { case "/ws", "/ws/": c.handleWebSocket(w, r) + case "/events", "/events/": + if r.Method == http.MethodGet { + c.handleSSE(w, r) + } else { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + } + case "/send", "/send/": + if r.Method == http.MethodPost { + c.handlePostSend(w, r) + } else { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + } default: http.NotFound(w, r) } @@ -297,16 +231,24 @@ func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error { msg.SessionID = sessionID var sent bool - for _, pc := range c.sessionConnectionsSnapshot(sessionID) { - if err := pc.writeJSON(msg); err != nil { - logger.DebugCF("pico", "Write to connection failed", map[string]any{ - "conn_id": pc.id, + c.subscribers.Range(func(key, value any) bool { + sub, ok := value.(picoSubscriber) + if !ok { + return 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 + }) if !sent { return fmt.Errorf("no active connections for session %s: %w", sessionID, channels.ErrSendFailed) @@ -314,6 +256,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() { @@ -327,12 +288,19 @@ 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 { maxConns = 100 } - if c.currentConnCount() >= maxConns { + if int(c.connCount.Load()) >= maxConns { http.Error(w, "too many connections", http.StatusServiceUnavailable) return } @@ -351,23 +319,14 @@ 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, err := c.createAndAddConnection(conn, sessionID, maxConns) - if err != nil { - _ = conn.WriteControl( - websocket.CloseMessage, - websocket.FormatCloseMessage(websocket.CloseTryAgainLater, "too many connections"), - time.Now().Add(2*time.Second), - ) - _ = conn.Close() - return + pc := &picoConn{ + id: uuid.New().String(), + conn: conn, + sessionID: sessionID, } + c.subscribers.Store(pc.id, pc) + c.connCount.Add(1) logger.InfoCF("pico", "WebSocket client connected", map[string]any{ "conn_id": pc.id, "session_id": sessionID, @@ -425,12 +384,13 @@ func (c *PicoChannel) matchedSubprotocol(r *http.Request) string { func (c *PicoChannel) readLoop(pc *picoConn) { defer func() { pc.close() - if removed := c.removeConnection(pc.id); removed != nil { - logger.InfoCF("pico", "WebSocket client disconnected", map[string]any{ - "conn_id": removed.id, - "session_id": removed.sessionID, - }) + 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, + }) }() readTimeout := time.Duration(c.config.ReadTimeout) * time.Second @@ -536,6 +496,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" @@ -544,7 +509,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{ @@ -562,7 +527,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..3df71b52f --- /dev/null +++ b/pkg/channels/pico/pico_http_test.go @@ -0,0 +1,214 @@ +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 newTestPicoHTTP(t *testing.T) *PicoChannel { + t.Helper() + cfg := config.PicoConfig{} + cfg.SetToken("tok") + ch, err := NewPicoChannel(cfg, bus.NewMessageBus()) + if err != nil { + t.Fatal(err) + } + return ch +} + +func TestPicoChannel_PostSend_Unauthorized(t *testing.T) { + t.Parallel() + ch := newTestPicoHTTP(t) + 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 := newTestPicoHTTP(t) + 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 := newTestPicoHTTP(t) + 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 := newTestPicoHTTP(t) + 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 := newTestPicoHTTP(t) + 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 := newTestPicoHTTP(t) + 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 := newTestPicoHTTP(t) + 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/pkg/channels/pico/pico_test.go b/pkg/channels/pico/pico_test.go index e712767ad..fe04c94e1 100644 --- a/pkg/channels/pico/pico_test.go +++ b/pkg/channels/pico/pico_test.go @@ -3,8 +3,6 @@ package pico import ( "context" "errors" - "fmt" - "sync" "testing" "github.com/sipeed/picoclaw/pkg/bus" @@ -26,92 +24,15 @@ func newTestPicoChannel(t *testing.T) *PicoChannel { return ch } -func TestCreateAndAddConnection_RespectsMaxConnectionsConcurrently(t *testing.T) { - ch := newTestPicoChannel(t) - - const ( - maxConns = 5 - goroutines = 64 - sessionID = "session-a" - ) - - var wg sync.WaitGroup - var mu sync.Mutex - successCount := 0 - errCount := 0 - - wg.Add(goroutines) - for i := 0; i < goroutines; i++ { - go func() { - defer wg.Done() - - pc, err := ch.createAndAddConnection(nil, sessionID, maxConns) - mu.Lock() - defer mu.Unlock() - - if err == nil { - successCount++ - if pc == nil { - t.Errorf("pc is nil on success") - } - return - } - if !errors.Is(err, channels.ErrTemporary) { - t.Errorf("unexpected error: %v", err) - return - } - errCount++ - }() - } - wg.Wait() - - if successCount > maxConns { - t.Fatalf("successCount=%d > maxConns=%d", successCount, maxConns) - } - if successCount+errCount != goroutines { - t.Fatalf("success=%d err=%d total=%d want=%d", successCount, errCount, successCount+errCount, goroutines) - } - if got := ch.currentConnCount(); got != maxConns { - t.Fatalf("currentConnCount=%d want=%d", got, maxConns) - } -} - -func TestRemoveConnection_CleansBothIndexes(t *testing.T) { - ch := newTestPicoChannel(t) - - pc, err := ch.createAndAddConnection(nil, "session-cleanup", 10) - if err != nil { - t.Fatalf("createAndAddConnection: %v", err) - } - - removed := ch.removeConnection(pc.id) - if removed == nil { - t.Fatal("removeConnection returned nil") - } - - ch.connsMu.RLock() - defer ch.connsMu.RUnlock() - - if _, ok := ch.connections[pc.id]; ok { - t.Fatalf("connID %s still exists in connections", pc.id) - } - if _, ok := ch.sessionConnections[pc.sessionID]; ok { - t.Fatalf("session %s still exists in sessionConnections", pc.sessionID) - } - if got := len(ch.connections); got != 0 { - t.Fatalf("len(connections)=%d want=0", got) - } -} - func TestBroadcastToSession_TargetsOnlyRequestedSession(t *testing.T) { ch := newTestPicoChannel(t) target := &picoConn{id: "target", sessionID: "s-target"} target.closed.Store(true) - ch.addConnForTest(target) + ch.addSubscriberForTest(target) other := &picoConn{id: "other", sessionID: "s-other"} - ch.addConnForTest(other) + ch.addSubscriberForTest(other) err := ch.broadcastToSession("pico:s-target", newMessage(TypeMessageCreate, map[string]any{"content": "hello"})) if err == nil { @@ -122,23 +43,7 @@ func TestBroadcastToSession_TargetsOnlyRequestedSession(t *testing.T) { } } -func (c *PicoChannel) addConnForTest(pc *picoConn) { - c.connsMu.Lock() - defer c.connsMu.Unlock() - if c.connections == nil { - c.connections = make(map[string]*picoConn) - } - if c.sessionConnections == nil { - c.sessionConnections = make(map[string]map[string]*picoConn) - } - if _, exists := c.connections[pc.id]; exists { - panic(fmt.Sprintf("duplicate conn id in test: %s", pc.id)) - } - c.connections[pc.id] = pc - bySession, ok := c.sessionConnections[pc.sessionID] - if !ok { - bySession = make(map[string]*picoConn) - c.sessionConnections[pc.sessionID] = bySession - } - bySession[pc.id] = pc +func (c *PicoChannel) addSubscriberForTest(sub picoSubscriber) { + c.subscribers.Store(sub.ID(), sub) + c.connCount.Add(1) } diff --git a/web/backend/api/pico.go b/web/backend/api/pico.go index 4faafc2ae..b60ad8cd9 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,