okooooooo/cc-connect
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add streaming preview and rate limiting support

Browse Source 
Streaming Preview:
- Real-time message preview during agent response (like "typing" indicator)
- Supported on Telegram, Discord, and Feishu
- Configurable interval, min delta chars, and max preview length

Rate Limiting:
- Per-session sliding-window rate limiter
- Prevents users from flooding the bot with messages
- Configurable max messages and window size

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
chenhg5andClaude Opus 4.6 committed 2026-03-06 07:02:37 +08:00
1 parent 00477f25a8
commit e5d328d9d6
12 files changed
+889 -8
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cmd/cc-connect/main.go
+37
View File
@@ -13,6 +13,7 @@ import (
"strconv"
"strings"
"syscall"
"time"
"github.com/chenhg5/cc-connect/config"
"github.com/chenhg5/cc-connect/core"
@@ -229,6 +230,42 @@ func main() {
}
engine.SetDisplayConfig(dcfg)
}
// Wire streaming preview
{
spcfg := core.DefaultStreamPreviewCfg()
if cfg.StreamPreview.Enabled != nil {
spcfg.Enabled = *cfg.StreamPreview.Enabled
}
if cfg.StreamPreview.IntervalMs != nil {
spcfg.IntervalMs = *cfg.StreamPreview.IntervalMs
}
if cfg.StreamPreview.MinDeltaChars != nil {
spcfg.MinDeltaChars = *cfg.StreamPreview.MinDeltaChars
}
if cfg.StreamPreview.MaxChars != nil {
spcfg.MaxChars = *cfg.StreamPreview.MaxChars
}
engine.SetStreamPreviewCfg(spcfg)
}
// Wire rate limiting
{
maxMsg := 20
windowSecs := 60
if cfg.RateLimit.MaxMessages != nil {
maxMsg = *cfg.RateLimit.MaxMessages
}
if cfg.RateLimit.WindowSecs != nil {
windowSecs = *cfg.RateLimit.WindowSecs
}
if maxMsg > 0 {
engine.SetRateLimitCfg(core.RateLimitCfg{
MaxMessages: maxMsg,
Window: time.Duration(windowSecs) * time.Second,
})
}
}
engine.SetDisplaySaveFunc(func(thinkingMaxLen, toolMaxLen *int) error {
return config.SaveDisplayConfig(thinkingMaxLen, toolMaxLen)
})
config.example.toml
+26
View File
@@ -45,6 +45,32 @@ level = "info" # debug, info, warn, error
# thinking_max_len = 300 # Max chars for thinking messages (default: 300) / 思考消息最大字符数（默认 300）
# tool_max_len = 500 # Max chars for tool use messages (default: 500) / 工具调用消息最大字符数（默认 500）
# =============================================================================
# Streaming Preview / 流式预览
# =============================================================================
# Real-time message preview: shows agent output as it streams, updating the
# message in-place (like a "typing" indicator). Supported on Telegram, Discord,
# and Feishu. The message is edited periodically until the final response is sent.
# 实时流式预览：Agent 输出时实时更新消息（类似"正在输入"效果）。
# 支持 Telegram、Discord 和飞书。消息会定期编辑，直到最终回复发送。
# [stream_preview]
# enabled = true # Enable/disable streaming preview (default: true) / 启用/禁用流式预览（默认 true）
# interval_ms = 1500 # Min ms between updates (default: 1500) / 更新最小间隔毫秒数（默认 1500）
# min_delta_chars = 30 # Min new chars before sending update (default: 30) / 发送更新前最少新增字符数（默认 30）
# max_chars = 2000 # Max preview length (default: 2000) / 预览最大长度（默认 2000）
# =============================================================================
# Rate Limiting / 速率限制
# =============================================================================
# Per-session sliding-window rate limiter. Prevents users from flooding the bot
# with too many messages. When exceeded, a friendly "slow down" message is sent.
# 每会话滑动窗口速率限制。防止用户刷消息。超限时返回友好的提示。
# [rate_limit]
# max_messages = 20 # Max messages per window; 0 = disabled (default: 20) / 窗口内最大消息数；0 = 禁用（默认 20）
# window_secs = 60 # Window size in seconds (default: 60) / 窗口时间秒数（默认 60）
# =============================================================================
# Speech-to-Text (Voice Messages) / 语音转文字（语音消息）
# =============================================================================
config/config.go
+18 -2
View File
@@ -24,8 +24,10 @@ type Config struct {
Log LogConfig `toml:"log"`
Language string `toml:"language"` // "en" or "zh", default is "en"
Speech SpeechConfig `toml:"speech"`
Display DisplayConfig `toml:"display"`
Quiet *bool `toml:"quiet,omitempty"` // global default for quiet mode; project-level overrides this
Display DisplayConfig `toml:"display"`
StreamPreview StreamPreviewConfig `toml:"stream_preview"` // real-time streaming preview
RateLimit RateLimitConfig `toml:"rate_limit"` // per-session rate limiting
Quiet *bool `toml:"quiet,omitempty"` // global default for quiet mode; project-level overrides this
}
// DisplayConfig controls how intermediate messages (thinking, tool output) are shown.
@@ -34,6 +36,20 @@ type DisplayConfig struct {
ToolMaxLen *int `toml:"tool_max_len"` // max chars for tool use messages; 0 = no truncation; default 500
}
// StreamPreviewConfig controls real-time streaming preview in IM.
type StreamPreviewConfig struct {
Enabled *bool `toml:"enabled"` // default true
IntervalMs *int `toml:"interval_ms"` // min ms between updates; default 1500
MinDeltaChars *int `toml:"min_delta_chars"` // min new chars before update; default 30
MaxChars *int `toml:"max_chars"` // max preview length; default 2000
}
// RateLimitConfig controls per-session message rate limiting.
type RateLimitConfig struct {
MaxMessages *int `toml:"max_messages"` // max messages per window; 0 = disabled; default 20
WindowSecs *int `toml:"window_secs"` // window size in seconds; default 60
}
// SpeechConfig configures speech-to-text for voice messages.
type SpeechConfig struct {
Enabled bool `toml:"enabled"`
core/engine.go
+53 -6
View File
@@ -105,6 +105,12 @@ type DisplayCfg struct {
ToolMaxLen int // max runes for tool use preview; 0 = no truncation
}
// RateLimitCfg controls per-session message rate limiting.
type RateLimitCfg struct {
MaxMessages int // max messages per window; 0 = disabled
Window time.Duration // sliding window size
}
// Engine routes messages between platforms and the agent for a single project.
type Engine struct {
name string
@@ -144,6 +150,9 @@ type Engine struct {
disabledCmds map[string]bool
rateLimiter *RateLimiter
streamPreview StreamPreviewCfg
// Interactive agent session management
interactiveMu sync.Mutex
interactiveStates map[string]*interactiveState // key = sessionKey
@@ -191,6 +200,7 @@ func NewEngine(name string, ag Agent, platforms []Platform, sessionStorePath str
aliases: make(map[string]string),
interactiveStates: make(map[string]*interactiveState),
startedAt: time.Now(),
streamPreview: DefaultStreamPreviewCfg(),
}
if cp, ok := ag.(CommandProvider); ok {
@@ -325,6 +335,16 @@ func (e *Engine) SetBannedWords(words []string) {
e.bannedWords = lower
}
// SetRateLimitCfg configures per-session message rate limiting.
func (e *Engine) SetRateLimitCfg(cfg RateLimitCfg) {
e.rateLimiter = NewRateLimiter(cfg.MaxMessages, cfg.Window)
}
// SetStreamPreviewCfg configures the streaming preview behavior.
func (e *Engine) SetStreamPreviewCfg(cfg StreamPreviewCfg) {
e.streamPreview = cfg
}
// RemoveCommand removes a custom command by name. Returns false if not found.
func (e *Engine) RemoveCommand(name string) bool {
return e.commands.Remove(name)
@@ -494,6 +514,13 @@ func (e *Engine) handleMessage(p Platform, msg *Message) {
content = e.resolveAlias(content)
msg.Content = content
// Rate limit check
if e.rateLimiter != nil && !e.rateLimiter.Allow(msg.SessionKey) {
slog.Info("message rate limited", "session", msg.SessionKey, "user", msg.UserName)
e.reply(p, msg.ReplyCtx, e.i18n.T(MsgRateLimited))
return
}
// Banned words check (skip for slash commands)
if !strings.HasPrefix(content, "/") {
if word := e.matchBannedWord(content); word != "" {
@@ -821,6 +848,10 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess
waitStart := time.Now()
firstEventLogged := false
state.mu.Lock()
sp := newStreamPreview(e.streamPreview, state.platform, state.replyCtx, e.ctx)
state.mu.Unlock()
for event := range state.agentSession.Events() {
if e.ctx.Err() != nil {
return
@@ -856,6 +887,9 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess
case EventText:
if event.Content != "" {
textParts = append(textParts, event.Content)
if sp.canPreview() {
sp.appendText(event.Content)
}
}
if event.SessionID != "" {
session.mu.Lock()
@@ -940,18 +974,26 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess
)
replyStart := time.Now()
for _, chunk := range splitMessage(fullResponse, maxPlatformMessageLen) {
if err := p.Send(e.ctx, replyCtx, chunk); err != nil {
slog.Error("failed to send reply", "error", err)
return
// If streaming preview was active, try to finalize in-place
if sp.finish(fullResponse) {
slog.Debug("stream preview: finalized in-place", "response_len", len(fullResponse))
} else {
for _, chunk := range splitMessage(fullResponse, maxPlatformMessageLen) {
if err := p.Send(e.ctx, replyCtx, chunk); err != nil {
slog.Error("failed to send reply", "error", err)
return
}
}
}
if elapsed := time.Since(replyStart); elapsed >= slowPlatformSend {
slog.Warn("slow final reply send", "platform", p.Name(), "elapsed", elapsed, "response_len", len(fullResponse))
}
return
case EventError:
sp.finish("") // clean up preview on error
if event.Error != nil {
slog.Error("agent error", "error", event.Error)
e.send(p, replyCtx, fmt.Sprintf(e.i18n.T(MsgError), event.Error))
@@ -972,8 +1014,13 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess
fullResponse := strings.Join(textParts, "")
session.AddHistory("assistant", fullResponse)
for _, chunk := range splitMessage(fullResponse, maxPlatformMessageLen) {
e.send(p, replyCtx, chunk)
if sp.finish(fullResponse) {
slog.Debug("stream preview: finalized in-place (process exited)")
} else {
for _, chunk := range splitMessage(fullResponse, maxPlatformMessageLen) {
e.send(p, replyCtx, chunk)
}
}
}
}
core/i18n.go
+8
View File
@@ -282,6 +282,7 @@ const (
MsgBannedWordBlocked MsgKey = "banned_word_blocked"
MsgCommandDisabled MsgKey = "command_disabled"
MsgRateLimited MsgKey = "rate_limited"
)
var messages = map[MsgKey]map[Language]string{
@@ -1537,6 +1538,13 @@ var messages = map[MsgKey]map[Language]string{
LangJapanese: "🚫 コマンド `%s` はこのプロジェクトで無効化されています。",
LangSpanish: "🚫 El comando `%s` está deshabilitado para este proyecto.",
},
MsgRateLimited: {
LangEnglish: "⏳ You are sending messages too fast. Please wait a moment.",
LangChinese: "⏳ 消息发送过快，请稍后再试。",
LangTraditionalChinese: "⏳ 訊息發送過快，請稍後再試。",
LangJapanese: "⏳ メッセージの送信が速すぎます。しばらくお待ちください。",
LangSpanish: "⏳ Estás enviando mensajes demasiado rápido. Espera un momento.",
},
}
func (i *I18n) T(key MsgKey) string {
core/ratelimit.go
+85
View File
@@ -0,0 +1,85 @@
package core
import (
"sync"
"time"
)
// RateLimiter implements a per-key sliding-window rate limiter.
// It tracks message timestamps per key and rejects requests that exceed
// the configured limit within the time window.
type RateLimiter struct {
mu sync.Mutex
buckets map[string]*rateBucket
maxMessages int
windowMs int64
}
type rateBucket struct {
timestamps []int64
lastAccess int64
}
// NewRateLimiter creates a rate limiter allowing maxMessages per window duration.
// Pass maxMessages=0 to disable rate limiting.
func NewRateLimiter(maxMessages int, window time.Duration) *RateLimiter {
rl := &RateLimiter{
buckets: make(map[string]*rateBucket),
maxMessages: maxMessages,
windowMs: window.Milliseconds(),
}
if maxMessages > 0 {
go rl.cleanupLoop()
}
return rl
}
// Allow checks whether a message from the given key is within the rate limit.
// Returns true if allowed (and records the timestamp), false if rate-limited.
func (rl *RateLimiter) Allow(key string) bool {
if rl.maxMessages <= 0 {
return true
}
rl.mu.Lock()
defer rl.mu.Unlock()
now := time.Now().UnixMilli()
b := rl.buckets[key]
if b == nil {
b = &rateBucket{}
rl.buckets[key] = b
}
b.lastAccess = now
cutoff := now - rl.windowMs
filtered := b.timestamps[:0]
for _, ts := range b.timestamps {
if ts > cutoff {
filtered = append(filtered, ts)
}
}
b.timestamps = filtered
if len(b.timestamps) >= rl.maxMessages {
return false
}
b.timestamps = append(b.timestamps, now)
return true
}
func (rl *RateLimiter) cleanupLoop() {
ticker := time.NewTicker(5 * time.Minute)
defer ticker.Stop()
for range ticker.C {
rl.mu.Lock()
now := time.Now().UnixMilli()
staleThreshold := rl.windowMs * 2
for k, b := range rl.buckets {
if now-b.lastAccess > staleThreshold {
delete(rl.buckets, k)
}
}
rl.mu.Unlock()
}
}
core/ratelimit_test.go
+77
View File
@@ -0,0 +1,77 @@
package core
import (
"sync"
"testing"
"time"
)
func TestRateLimiter_AllowWithinLimit(t *testing.T) {
rl := NewRateLimiter(5, time.Minute)
for i := 0; i < 5; i++ {
if !rl.Allow("user1") {
t.Errorf("request %d should be allowed", i+1)
}
}
}
func TestRateLimiter_BlockExceedingLimit(t *testing.T) {
rl := NewRateLimiter(3, time.Minute)
for i := 0; i < 3; i++ {
rl.Allow("user1")
}
if rl.Allow("user1") {
t.Error("4th request should be blocked")
}
}
func TestRateLimiter_DifferentKeys(t *testing.T) {
rl := NewRateLimiter(2, time.Minute)
rl.Allow("user1")
rl.Allow("user1")
if rl.Allow("user1") {
t.Error("user1 should be blocked")
}
if !rl.Allow("user2") {
t.Error("user2 should be allowed (independent bucket)")
}
}
func TestRateLimiter_WindowExpiry(t *testing.T) {
rl := NewRateLimiter(2, 50*time.Millisecond)
rl.Allow("user1")
rl.Allow("user1")
if rl.Allow("user1") {
t.Error("should be blocked immediately")
}
time.Sleep(60 * time.Millisecond)
if !rl.Allow("user1") {
t.Error("should be allowed after window expires")
}
}
func TestRateLimiter_Disabled(t *testing.T) {
rl := NewRateLimiter(0, time.Minute)
for i := 0; i < 100; i++ {
if !rl.Allow("user1") {
t.Error("should always allow when disabled")
}
}
}
func TestRateLimiter_Concurrent(t *testing.T) {
rl := NewRateLimiter(100, time.Minute)
var wg sync.WaitGroup
for i := 0; i < 200; i++ {
wg.Add(1)
go func() {
defer wg.Done()
rl.Allow("user1")
}()
}
wg.Wait()
}
core/streaming.go
+241
View File
@@ -0,0 +1,241 @@
package core
import (
"context"
"log/slog"
"sync"
"time"
)
// StreamPreviewCfg controls the streaming preview behavior.
type StreamPreviewCfg struct {
Enabled bool // global toggle
IntervalMs int // minimum ms between updates (default 1500)
MinDeltaChars int // minimum new chars before sending an update (default 30)
MaxChars int // max preview length (default 2000)
}
// DefaultStreamPreviewCfg returns sensible defaults.
func DefaultStreamPreviewCfg() StreamPreviewCfg {
return StreamPreviewCfg{
Enabled: true,
IntervalMs: 1500,
MinDeltaChars: 30,
MaxChars: 2000,
}
}
// streamPreview manages the state and throttling of a single streaming preview.
// It accumulates text from EventText events and periodically pushes
// updates to the platform via MessageUpdater.UpdateMessage.
type streamPreview struct {
mu sync.Mutex
cfg StreamPreviewCfg
platform Platform
replyCtx any
ctx context.Context
fullText string // accumulated full text so far
lastSentText string // what was last successfully sent to the platform
lastSentAt time.Time
previewMsgID any // platform-specific ID for the preview message (returned by SendPreviewStart)
degraded bool // if true, stop trying (platform doesn't support it or permanent error)
timer *time.Timer
timerStop chan struct{} // closed when preview ends
}
// PreviewStarter is an optional interface for platforms that can initiate a
// streaming preview message and return a handle for subsequent updates.
type PreviewStarter interface {
// SendPreviewStart sends the initial preview message and returns a handle
// that can be passed to UpdateMessage for edits. Returns nil handle if
// preview is not supported for this context.
SendPreviewStart(ctx context.Context, replyCtx any, content string) (previewHandle any, err error)
}
// PreviewCleaner is an optional interface for platforms that need to clean up
// the preview message after the final response is sent (e.g. Discord deletes
// the preview and sends a fresh message).
type PreviewCleaner interface {
DeletePreviewMessage(ctx context.Context, previewHandle any) error
}
func newStreamPreview(cfg StreamPreviewCfg, p Platform, replyCtx any, ctx context.Context) *streamPreview {
return &streamPreview{
cfg: cfg,
platform: p,
replyCtx: replyCtx,
ctx: ctx,
timerStop: make(chan struct{}),
}
}
// canPreview returns true if the platform supports message updating.
func (sp *streamPreview) canPreview() bool {
if sp.degraded || !sp.cfg.Enabled {
return false
}
_, ok := sp.platform.(MessageUpdater)
return ok
}
// appendText adds new text content and triggers a throttled flush if needed.
func (sp *streamPreview) appendText(text string) {
sp.mu.Lock()
defer sp.mu.Unlock()
if sp.degraded || !sp.cfg.Enabled {
return
}
sp.fullText += text
displayText := sp.fullText
maxChars := sp.cfg.MaxChars
if maxChars > 0 && len([]rune(displayText)) > maxChars {
displayText = string([]rune(displayText)[:maxChars]) + "…"
}
delta := len([]rune(displayText)) - len([]rune(sp.lastSentText))
elapsed := time.Since(sp.lastSentAt)
interval := time.Duration(sp.cfg.IntervalMs) * time.Millisecond
if delta < sp.cfg.MinDeltaChars && !sp.lastSentAt.IsZero() {
sp.scheduleFlushLocked(interval)
return
}
if elapsed < interval && !sp.lastSentAt.IsZero() {
remaining := interval - elapsed
sp.scheduleFlushLocked(remaining)
return
}
sp.cancelTimerLocked()
sp.flushLocked(displayText)
}
func (sp *streamPreview) scheduleFlushLocked(delay time.Duration) {
if sp.timer != nil {
return // already scheduled
}
sp.timer = time.AfterFunc(delay, func() {
sp.mu.Lock()
defer sp.mu.Unlock()
sp.timer = nil
if sp.degraded {
return
}
displayText := sp.fullText
maxChars := sp.cfg.MaxChars
if maxChars > 0 && len([]rune(displayText)) > maxChars {
displayText = string([]rune(displayText)[:maxChars]) + "…"
}
sp.flushLocked(displayText)
})
}
func (sp *streamPreview) cancelTimerLocked() {
if sp.timer != nil {
sp.timer.Stop()
sp.timer = nil
}
}
// flushLocked sends the current preview text to the platform. Must hold sp.mu.
func (sp *streamPreview) flushLocked(text string) {
if text == sp.lastSentText || text == "" {
return
}
updater, ok := sp.platform.(MessageUpdater)
if !ok {
sp.degraded = true
return
}
if sp.previewMsgID == nil {
// First preview: try to send a new preview message
if starter, ok := sp.platform.(PreviewStarter); ok {
handle, err := starter.SendPreviewStart(sp.ctx, sp.replyCtx, text+"▍")
if err != nil {
slog.Debug("stream preview: start failed, degrading", "error", err)
sp.degraded = true
return
}
sp.previewMsgID = handle
} else {
// Platform supports UpdateMessage but not PreviewStarter;
// use Send to create initial message, then update in-place
if err := sp.platform.Send(sp.ctx, sp.replyCtx, text+"▍"); err != nil {
slog.Debug("stream preview: initial send failed", "error", err)
sp.degraded = true
return
}
// For platforms without PreviewStarter, replyCtx itself serves as the handle
sp.previewMsgID = sp.replyCtx
}
sp.lastSentText = text
sp.lastSentAt = time.Now()
return
}
// Update existing preview message
if err := updater.UpdateMessage(sp.ctx, sp.previewMsgID, text+"▍"); err != nil {
slog.Debug("stream preview: update failed, degrading", "error", err)
sp.degraded = true
return
}
sp.lastSentText = text
sp.lastSentAt = time.Now()
}
// finish is called when the agent response is complete. It cancels any pending
// timer and optionally cleans up the preview message.
// Returns true if a preview was active and the final message was sent via preview
// (so the caller should skip sending the full response separately).
func (sp *streamPreview) finish(finalText string) bool {
sp.mu.Lock()
defer sp.mu.Unlock()
sp.cancelTimerLocked()
close(sp.timerStop)
if sp.previewMsgID == nil || sp.degraded {
return false
}
// If platform wants to delete the preview and send fresh, let it
if cleaner, ok := sp.platform.(PreviewCleaner); ok {
_ = cleaner.DeletePreviewMessage(sp.ctx, sp.previewMsgID)
return false // caller should send the final message normally
}
// Otherwise update the preview message in-place with the final text
updater, ok := sp.platform.(MessageUpdater)
if !ok {
return false
}
// For very long responses, we may need chunked sending instead
maxChars := sp.cfg.MaxChars
if maxChars > 0 && len([]rune(finalText)) > maxChars {
// Final text exceeds preview limit; delete preview and let caller handle it
return false
}
if err := updater.UpdateMessage(sp.ctx, sp.previewMsgID, finalText); err != nil {
slog.Debug("stream preview: final update failed", "error", err)
return false
}
return true
}
// getFullText returns the accumulated text so far.
func (sp *streamPreview) getFullText() string {
sp.mu.Lock()
defer sp.mu.Unlock()
return sp.fullText
}
core/streaming_test.go
+184
View File
@@ -0,0 +1,184 @@
package core
import (
"context"
"sync"
"testing"
"time"
)
// mockUpdaterPlatform implements Platform + MessageUpdater + PreviewStarter.
type mockUpdaterPlatform struct {
stubPlatformEngine
mu sync.Mutex
messages []string // track all sent/updated messages
lastMsg string
}
func (m *mockUpdaterPlatform) SendPreviewStart(_ context.Context, _ any, content string) (any, error) {
m.mu.Lock()
defer m.mu.Unlock()
m.messages = append(m.messages, "start:"+content)
m.lastMsg = content
return "preview-handle", nil
}
func (m *mockUpdaterPlatform) UpdateMessage(_ context.Context, _ any, content string) error {
m.mu.Lock()
defer m.mu.Unlock()
m.messages = append(m.messages, "update:"+content)
m.lastMsg = content
return nil
}
func (m *mockUpdaterPlatform) getMessages() []string {
m.mu.Lock()
defer m.mu.Unlock()
out := make([]string, len(m.messages))
copy(out, m.messages)
return out
}
func TestStreamPreview_BasicFlow(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 100,
MinDeltaChars: 5,
MaxChars: 500,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
if !sp.canPreview() {
t.Fatal("should be able to preview")
}
sp.appendText("Hello ")
time.Sleep(150 * time.Millisecond)
msgs := mp.getMessages()
if len(msgs) == 0 {
t.Fatal("expected at least one message sent")
}
if msgs[0] != "start:Hello ▍" {
t.Errorf("first message = %q, want 'start:Hello ▍'", msgs[0])
}
}
func TestStreamPreview_ThrottlesUpdates(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 200,
MinDeltaChars: 5,
MaxChars: 500,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
// Rapid-fire small appends
for i := 0; i < 10; i++ {
sp.appendText("ab")
time.Sleep(10 * time.Millisecond)
}
// Wait for throttle timers to fire
time.Sleep(300 * time.Millisecond)
msgs := mp.getMessages()
// Should NOT have 10 individual updates; throttling should batch them
if len(msgs) >= 10 {
t.Errorf("expected throttling to reduce updates, got %d", len(msgs))
}
if len(msgs) == 0 {
t.Error("expected at least one update")
}
}
func TestStreamPreview_MaxChars(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 50,
MinDeltaChars: 1,
MaxChars: 10,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
sp.appendText("This is a very long text that exceeds max chars limit")
time.Sleep(100 * time.Millisecond)
msgs := mp.getMessages()
if len(msgs) == 0 {
t.Fatal("expected at least one message")
}
// Last message should be truncated
for _, m := range msgs {
if len(m) > 0 {
// Content after "start:" or "update:" should respect maxChars
content := m
for _, prefix := range []string{"start:", "update:"} {
if len(content) > len(prefix) && content[:len(prefix)] == prefix {
content = content[len(prefix):]
}
}
if len([]rune(content)) > 15 { // 10 chars + "…" + "▍" with some margin
t.Errorf("message too long: %q (%d runes)", content, len([]rune(content)))
}
}
}
}
func TestStreamPreview_Disabled(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{Enabled: false}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
if sp.canPreview() {
t.Error("should not be able to preview when disabled")
}
sp.appendText("Hello")
time.Sleep(50 * time.Millisecond)
msgs := mp.getMessages()
if len(msgs) != 0 {
t.Error("no messages should be sent when disabled")
}
}
func TestStreamPreview_FinishInPlace(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 50,
MinDeltaChars: 1,
MaxChars: 500,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
sp.appendText("Hello World")
time.Sleep(100 * time.Millisecond)
ok := sp.finish("Hello World Final")
if !ok {
t.Error("finish should return true when preview was active")
}
msgs := mp.getMessages()
last := msgs[len(msgs)-1]
if last != "update:Hello World Final" {
t.Errorf("last message = %q, want 'update:Hello World Final'", last)
}
}
func TestStreamPreview_NonUpdaterPlatform(t *testing.T) {
p := &stubPlatformEngine{n: "plain"}
cfg := DefaultStreamPreviewCfg()
sp := newStreamPreview(cfg, p, "ctx", context.Background())
if sp.canPreview() {
t.Error("should not preview on non-updater platform")
}
}
platform/discord/discord.go
+54
View File
@@ -455,6 +455,60 @@ func (p *Platform) ReconstructReplyCtx(sessionKey string) (any, error) {
return replyContext{channelID: parts[1]}, nil
}
// discordPreviewHandle stores the IDs needed to edit or delete a preview message.
type discordPreviewHandle struct {
channelID string
messageID string
}
// SendPreviewStart sends a new message and returns a handle for subsequent edits.
func (p *Platform) SendPreviewStart(ctx context.Context, rctx any, content string) (any, error) {
var channelID string
switch rc := rctx.(type) {
case replyContext:
channelID = rc.channelID
case *interactionReplyCtx:
channelID = rc.channelID
default:
return nil, fmt.Errorf("discord: invalid reply context type %T", rctx)
}
if len(content) > maxDiscordLen {
content = content[:maxDiscordLen]
}
sent, err := p.session.ChannelMessageSend(channelID, content)
if err != nil {
return nil, fmt.Errorf("discord: send preview: %w", err)
}
return &discordPreviewHandle{channelID: channelID, messageID: sent.ID}, nil
}
// UpdateMessage edits an existing message identified by previewHandle.
func (p *Platform) UpdateMessage(ctx context.Context, previewHandle any, content string) error {
h, ok := previewHandle.(*discordPreviewHandle)
if !ok {
return fmt.Errorf("discord: invalid preview handle type %T", previewHandle)
}
if len(content) > maxDiscordLen {
content = content[:maxDiscordLen]
}
_, err := p.session.ChannelMessageEdit(h.channelID, h.messageID, content)
if err != nil {
return fmt.Errorf("discord: edit message: %w", err)
}
return nil
}
// DeletePreviewMessage removes the preview message so the final response can
// be sent as a fresh message (avoids notification confusion).
func (p *Platform) DeletePreviewMessage(ctx context.Context, previewHandle any) error {
h, ok := previewHandle.(*discordPreviewHandle)
if !ok {
return fmt.Errorf("discord: invalid preview handle type %T", previewHandle)
}
return p.session.ChannelMessageDelete(h.channelID, h.messageID)
}
func (p *Platform) Stop() error {
if p.session != nil {
return p.session.Close()
platform/feishu/feishu.go
+70
View File
@@ -607,6 +607,76 @@ func (p *Platform) ReconstructReplyCtx(sessionKey string) (any, error) {
return replyContext{chatID: parts[1]}, nil
}
// feishuPreviewHandle stores the message ID for an editable preview message.
type feishuPreviewHandle struct {
messageID string
chatID string
}
// SendPreviewStart sends a new text message and returns a handle for subsequent edits.
func (p *Platform) SendPreviewStart(ctx context.Context, rctx any, content string) (any, error) {
rc, ok := rctx.(replyContext)
if !ok {
return nil, fmt.Errorf("feishu: invalid reply context type %T", rctx)
}
chatID := rc.chatID
if chatID == "" {
return nil, fmt.Errorf("feishu: chatID is empty")
}
b, _ := json.Marshal(map[string]string{"text": content})
resp, err := p.client.Im.Message.Create(ctx, larkim.NewCreateMessageReqBuilder().
ReceiveIdType(larkim.ReceiveIdTypeChatId).
Body(larkim.NewCreateMessageReqBodyBuilder().
ReceiveId(chatID).
MsgType(larkim.MsgTypeText).
Content(string(b)).
Build()).
Build())
if err != nil {
return nil, fmt.Errorf("feishu: send preview: %w", err)
}
if !resp.Success() {
return nil, fmt.Errorf("feishu: send preview code=%d msg=%s", resp.Code, resp.Msg)
}
msgID := ""
if resp.Data != nil && resp.Data.MessageId != nil {
msgID = *resp.Data.MessageId
}
if msgID == "" {
return nil, fmt.Errorf("feishu: send preview: no message ID returned")
}
return &feishuPreviewHandle{messageID: msgID, chatID: chatID}, nil
}
// UpdateMessage edits an existing text message identified by previewHandle.
func (p *Platform) UpdateMessage(ctx context.Context, previewHandle any, content string) error {
h, ok := previewHandle.(*feishuPreviewHandle)
if !ok {
return fmt.Errorf("feishu: invalid preview handle type %T", previewHandle)
}
b, _ := json.Marshal(map[string]string{"text": content})
msgType := larkim.MsgTypeText
resp, err := p.client.Im.Message.Update(ctx, larkim.NewUpdateMessageReqBuilder().
MessageId(h.messageID).
Body(larkim.NewUpdateMessageReqBodyBuilder().
MsgType(msgType).
Content(string(b)).
Build()).
Build())
if err != nil {
return fmt.Errorf("feishu: update message: %w", err)
}
if !resp.Success() {
return fmt.Errorf("feishu: update message code=%d msg=%s", resp.Code, resp.Msg)
}
return nil
}
func (p *Platform) Stop() error {
if p.cancel != nil {
p.cancel()
platform/telegram/telegram.go
+36
View File
@@ -464,6 +464,42 @@ func (p *Platform) ReconstructReplyCtx(sessionKey string) (any, error) {
return replyContext{chatID: chatID}, nil
}
// telegramPreviewHandle stores the chat and message IDs for an editable preview message.
type telegramPreviewHandle struct {
chatID int64
messageID int
}
// SendPreviewStart sends a new message and returns a handle for subsequent edits.
func (p *Platform) SendPreviewStart(ctx context.Context, rctx any, content string) (any, error) {
rc, ok := rctx.(replyContext)
if !ok {
return nil, fmt.Errorf("telegram: invalid reply context type %T", rctx)
}
msg := tgbotapi.NewMessage(rc.chatID, content)
msg.ParseMode = ""
sent, err := p.bot.Send(msg)
if err != nil {
return nil, fmt.Errorf("telegram: send preview: %w", err)
}
return &telegramPreviewHandle{chatID: rc.chatID, messageID: sent.MessageID}, nil
}
// UpdateMessage edits an existing message identified by previewHandle.
func (p *Platform) UpdateMessage(ctx context.Context, previewHandle any, content string) error {
h, ok := previewHandle.(*telegramPreviewHandle)
if !ok {
return fmt.Errorf("telegram: invalid preview handle type %T", previewHandle)
}
edit := tgbotapi.NewEditMessageText(h.chatID, h.messageID, content)
_, err := p.bot.Send(edit)
if err != nil {
return fmt.Errorf("telegram: edit message: %w", err)
}
return nil
}
func (p *Platform) Stop() error {
if p.cancel != nil {
p.cancel()