diff --git a/examples/pico-echo-server/main.go b/examples/pico-echo-server/main.go index 69d55462d..46970fb34 100644 --- a/examples/pico-echo-server/main.go +++ b/examples/pico-echo-server/main.go @@ -43,8 +43,7 @@ type server struct { func (s *server) handleWS(w http.ResponseWriter, r *http.Request) { if s.token != "" { auth := r.Header.Get("Authorization") - qtoken := r.URL.Query().Get("token") - if auth != "Bearer "+s.token && qtoken != s.token { + if auth != "Bearer "+s.token { http.Error(w, "unauthorized", http.StatusUnauthorized) return } @@ -141,22 +140,21 @@ func main() { http.HandleFunc("/ws", s.handleWS) + log.Printf("listening on %s", *addr) + log.Printf("connect with: ws://localhost%s/ws", *addr) + fmt.Println("Type messages to send to connected clients (Ctrl+C to quit):") + go func() { - log.Printf("listening on %s (token=%q)", *addr, *token) - log.Printf("connect with: ws://localhost%s/ws", *addr) - if err := http.ListenAndServe(*addr, nil); err != nil { - log.Fatal(err) + scanner := bufio.NewScanner(os.Stdin) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + s.broadcast(line) + log.Printf("[server] sent: %s", line) } }() - fmt.Println("Type messages to send to connected clients (Ctrl+C to quit):") - scanner := bufio.NewScanner(os.Stdin) - for scanner.Scan() { - line := strings.TrimSpace(scanner.Text()) - if line == "" { - continue - } - s.broadcast(line) - log.Printf("[server] sent: %s", line) - } + log.Fatal(http.ListenAndServe(*addr, nil)) } diff --git a/pkg/channels/pico/client.go b/pkg/channels/pico/client.go index f03b2a200..2c335050d 100644 --- a/pkg/channels/pico/client.go +++ b/pkg/channels/pico/client.go @@ -93,10 +93,13 @@ func (c *PicoClientChannel) dial() error { return err } + connCtx, connCancel := context.WithCancel(c.ctx) + pc := &picoConn{ id: uuid.New().String(), conn: ws, sessionID: c.config.SessionID, + cancel: connCancel, } if pc.sessionID == "" { pc.sessionID = uuid.New().String() @@ -106,7 +109,7 @@ func (c *PicoClientChannel) dial() error { c.conn = pc c.mu.Unlock() - go c.readLoop(pc) + go c.readLoop(connCtx, pc) return nil } @@ -148,7 +151,7 @@ func (c *PicoClientChannel) reconnectLoop() { } } -func (c *PicoClientChannel) readLoop(pc *picoConn) { +func (c *PicoClientChannel) readLoop(connCtx context.Context, pc *picoConn) { defer pc.close() readTimeout := time.Duration(c.config.ReadTimeout) * time.Second @@ -165,11 +168,11 @@ func (c *PicoClientChannel) readLoop(pc *picoConn) { if pingInterval <= 0 { pingInterval = 30 * time.Second } - go c.pingLoop(pc, pingInterval) + go c.pingLoop(connCtx, pc, pingInterval) for { select { - case <-c.ctx.Done(): + case <-connCtx.Done(): return default: } @@ -199,12 +202,12 @@ func (c *PicoClientChannel) readLoop(pc *picoConn) { } } -func (c *PicoClientChannel) pingLoop(pc *picoConn, interval time.Duration) { +func (c *PicoClientChannel) pingLoop(connCtx context.Context, pc *picoConn, interval time.Duration) { ticker := time.NewTicker(interval) defer ticker.Stop() for { select { - case <-c.ctx.Done(): + case <-connCtx.Done(): return case <-ticker.C: if pc.closed.Load() { @@ -303,8 +306,14 @@ func (c *PicoClientChannel) StartTyping(ctx context.Context, chatID string) (fun return func() {}, err } return func() { + c.mu.Lock() + currentPC := c.conn + c.mu.Unlock() + if currentPC == nil { + return + } stopMsg := newMessage(TypeTypingStop, nil) stopMsg.SessionID = strings.TrimPrefix(chatID, "pico_client:") - pc.writeJSON(stopMsg) + currentPC.writeJSON(stopMsg) }, nil } diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index 8d8b62a67..600eaaf13 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -27,6 +27,7 @@ type picoConn struct { sessionID string writeMu sync.Mutex closed atomic.Bool + cancel context.CancelFunc // cancels per-connection goroutines (e.g. pingLoop) } // writeJSON sends a JSON message to the connection with write locking. @@ -42,6 +43,9 @@ func (pc *picoConn) writeJSON(v any) error { // close closes the connection. func (pc *picoConn) close() { if pc.closed.CompareAndSwap(false, true) { + if pc.cancel != nil { + pc.cancel() + } pc.conn.Close() } }