diff --git a/cmd/cc-connect/main.go b/cmd/cc-connect/main.go index 47da323..0011acd 100644 --- a/cmd/cc-connect/main.go +++ b/cmd/cc-connect/main.go @@ -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) }) diff --git a/config.example.toml b/config.example.toml index 803c78f..0a8f239 100644 --- a/config.example.toml +++ b/config.example.toml @@ -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) / 语音转文字（语音消息） # ============================================================================= diff --git a/config/config.go b/config/config.go index 88416a3..3b390f1 100644 --- a/config/config.go +++ b/config/config.go @@ -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"` diff --git a/core/engine.go b/core/engine.go index 9f9de06..60fc81b 100644 --- a/core/engine.go +++ b/core/engine.go @@ -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) + } } } } diff --git a/core/i18n.go b/core/i18n.go index 32132ff..1c24676 100644 --- a/core/i18n.go +++ b/core/i18n.go @@ -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 { diff --git a/core/ratelimit.go b/core/ratelimit.go new file mode 100644 index 0000000..37749a8 --- /dev/null +++ b/core/ratelimit.go @@ -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() + } +} diff --git a/core/ratelimit_test.go b/core/ratelimit_test.go new file mode 100644 index 0000000..b83570e --- /dev/null +++ b/core/ratelimit_test.go @@ -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() +} diff --git a/core/streaming.go b/core/streaming.go new file mode 100644 index 0000000..40624e6 --- /dev/null +++ b/core/streaming.go @@ -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 +} diff --git a/core/streaming_test.go b/core/streaming_test.go new file mode 100644 index 0000000..131cdf6 --- /dev/null +++ b/core/streaming_test.go @@ -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") + } +} diff --git a/platform/discord/discord.go b/platform/discord/discord.go index b3a0ce1..7d79630 100644 --- a/platform/discord/discord.go +++ b/platform/discord/discord.go @@ -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() diff --git a/platform/feishu/feishu.go b/platform/feishu/feishu.go index 0827db7..03e9df6 100644 --- a/platform/feishu/feishu.go +++ b/platform/feishu/feishu.go @@ -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() diff --git a/platform/telegram/telegram.go b/platform/telegram/telegram.go index b6b7076..5b97abd 100644 --- a/platform/telegram/telegram.go +++ b/platform/telegram/telegram.go @@ -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()