fix: address code review feedback on streaming PR

- Deduplicate Streamer interface: alias channels.Streamer to bus.Streamer
  to prevent type drift across packages
- Increase SSE scanner buffer to 10MB max to handle large single-line
  responses that exceed bufio.Scanner's 64KB default
- Switch draftID generation from math/rand to crypto/rand for
  collision-resistant random IDs
- Add context cancellation check in SSE parsing loop so cancelled
  streams stop processing immediately
- Log Finalize failures with chat_id and content length for debugging
  silent message delivery failures
This commit is contained in:
Amir Mamaghani 2026-03-05 18:45:42 +01:00
parent 9f3caa3d6c
commit fa3c00ad9c
3 changed files with 31 additions and 14 deletions

View file

@ -1,6 +1,10 @@
package channels
import "context"
import (
"context"
"github.com/sipeed/picoclaw/pkg/bus"
)
// TypingCapable — channels that can show a typing/thinking indicator.
// StartTyping begins the indicator and returns a stop function.
@ -44,15 +48,9 @@ type StreamingCapable interface {
BeginStream(ctx context.Context, chatID string) (Streamer, error)
}
// Streamer pushes incremental content to a streaming-capable channel.
type Streamer interface {
// Update sends accumulated partial content to the user.
Update(ctx context.Context, content string) error
// Finalize commits the final message. After this, the Streamer is done.
Finalize(ctx context.Context, content string) error
// Cancel aborts the stream (e.g. when tool calls are detected mid-stream).
Cancel(ctx context.Context)
}
// Streamer is defined in pkg/bus to avoid circular imports.
// This alias keeps channel implementations using channels.Streamer unchanged.
type Streamer = bus.Streamer
// PlaceholderRecorder is injected into channels by Manager.
// Channels call these methods on inbound to register typing/placeholder state.

View file

@ -3,7 +3,8 @@ package telegram
import (
"context"
"fmt"
"math/rand"
"crypto/rand"
"encoding/binary"
"net/http"
"net/url"
"os"
@ -800,7 +801,7 @@ func (c *TelegramChannel) BeginStream(ctx context.Context, chatID string) (chann
return &telegramStreamer{
bot: c.bot,
chatID: cid,
draftID: rand.Intn(1<<31-1) + 1, // non-zero random draft ID
draftID: cryptoRandInt(), // non-zero random draft ID
}, nil
}
@ -868,6 +869,11 @@ func (s *telegramStreamer) Finalize(ctx context.Context, content string) error {
// Fallback to plain text
tgMsg.ParseMode = ""
if _, err = s.bot.SendMessage(ctx, tgMsg); err != nil {
logger.ErrorCF("telegram", "Finalize failed after HTML and plain-text attempts", map[string]any{
"chat_id": s.chatID,
"error": err.Error(),
"len": len(content),
})
return fmt.Errorf("telegram finalize: %w", err)
}
}
@ -877,3 +883,10 @@ func (s *telegramStreamer) Finalize(ctx context.Context, content string) error {
func (s *telegramStreamer) Cancel(ctx context.Context) {
// Draft auto-expires on Telegram's side; nothing to clean up.
}
// cryptoRandInt returns a non-zero random int using crypto/rand.
func cryptoRandInt() int {
var b [4]byte
_, _ = rand.Read(b[:])
return int(binary.BigEndian.Uint32(b[:])) | 1 // ensure non-zero
}

View file

@ -244,11 +244,11 @@ func (p *Provider) ChatStream(
return nil, fmt.Errorf("API request failed:\n Status: %d\n Body: %s", resp.StatusCode, string(body))
}
return parseStreamResponse(resp.Body, onChunk)
return parseStreamResponse(ctx, resp.Body, onChunk)
}
// parseStreamResponse parses an OpenAI-compatible SSE stream.
func parseStreamResponse(reader io.Reader, onChunk func(accumulated string)) (*LLMResponse, error) {
func parseStreamResponse(ctx context.Context, reader io.Reader, onChunk func(accumulated string)) (*LLMResponse, error) {
var textContent strings.Builder
var finishReason string
var usage *UsageInfo
@ -262,7 +262,13 @@ func parseStreamResponse(reader io.Reader, onChunk func(accumulated string)) (*L
activeTools := map[int]*toolAccum{}
scanner := bufio.NewScanner(reader)
scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024) // 1MB initial, 10MB max
for scanner.Scan() {
// Check for context cancellation between chunks
if err := ctx.Err(); err != nil {
return nil, err
}
line := scanner.Text()
if !strings.HasPrefix(line, "data: ") {