From ce5957c6c0e08f92b133bbced39009799447737a Mon Sep 17 00:00:00 2001 From: "mudler's LocalAI [bot]" <139863280+localai-bot@users.noreply.github.com> Date: Sat, 22 Aug 2026 10:11:51 +0200 Subject: [PATCH] fix(telegram): preserve progressive drafts (#489) Agent cleanup canceled the preview context before the connector drained its queued deltas. Keep the consumer alive for the whole Telegram turn. Use the standard sendMessageDraft transport for private chats and retain rich Markdown for final responses. This matches the Hermes streaming lifecycle and visible cursor behavior. Co-authored-by: Ettore Di Giacinto --- services/connectors/telegram.go | 5 +- services/connectors/telegram_api.go | 16 +++--- services/connectors/telegram_api_test.go | 30 +++++------- .../connectors/telegram_integration_test.go | 49 ++++++++++++------- services/connectors/telegram_stream.go | 18 +++++-- services/connectors/telegram_stream_test.go | 38 +++++++------- 6 files changed, 88 insertions(+), 68 deletions(-) diff --git a/services/connectors/telegram.go b/services/connectors/telegram.go index 9a2b883..a734c31 100644 --- a/services/connectors/telegram.go +++ b/services/connectors/telegram.go @@ -76,7 +76,10 @@ func telegramNewJobWithStream(parent context.Context, api telegramAPI, chatID in metadata[telegramStreamingMetadataKey] = true opts := append(telegramAskOptions(history, jobUUID, metadata, nil), types.WithContext(parent)) job := types.NewJob(opts...) - session := newTelegramStreamSessionWithContexts(parent, job.GetContext(), api, chatID, private, delivery, telegramDraftHeartbeatInterval) + // Like Hermes' StreamConsumer, the preview queue belongs to the whole + // connector turn. Agent.consumeJob cancels the job context during normal + // completion, before this handler has had a chance to flush and finalize. + session := newTelegramStreamSession(parent, api, chatID, private, delivery) job.StreamCallback = session.Accept return job, session } diff --git a/services/connectors/telegram_api.go b/services/connectors/telegram_api.go index f804d70..2407ed2 100644 --- a/services/connectors/telegram_api.go +++ b/services/connectors/telegram_api.go @@ -13,7 +13,7 @@ import ( const telegramBotAPIBaseURL = "https://api.telegram.org" type telegramAPI interface { - sendRichMessageDraft(context.Context, telegramRichMessageDraft) error + sendMessageDraft(context.Context, telegramMessageDraft) error sendRichMessage(context.Context, telegramRichMessage) error } @@ -21,10 +21,10 @@ type telegramInputRichMessage struct { Markdown string `json:"markdown"` } -type telegramRichMessageDraft struct { - ChatID int64 `json:"chat_id"` - DraftID int64 `json:"draft_id"` - RichMessage telegramInputRichMessage `json:"rich_message"` +type telegramMessageDraft struct { + ChatID int64 `json:"chat_id"` + DraftID int64 `json:"draft_id"` + Text string `json:"text"` } type telegramRichMessage struct { @@ -68,11 +68,11 @@ func newTelegramHTTPAPI(token string, client *http.Client, baseURL string) teleg } } -func (a *telegramHTTPAPI) sendRichMessageDraft(ctx context.Context, input telegramRichMessageDraft) error { +func (a *telegramHTTPAPI) sendMessageDraft(ctx context.Context, input telegramMessageDraft) error { if input.DraftID == 0 { - return fmt.Errorf("telegram sendRichMessageDraft: draft_id must be nonzero") + return fmt.Errorf("telegram sendMessageDraft: draft_id must be nonzero") } - return a.call(ctx, "sendRichMessageDraft", input) + return a.call(ctx, "sendMessageDraft", input) } func (a *telegramHTTPAPI) sendRichMessage(ctx context.Context, input telegramRichMessage) error { diff --git a/services/connectors/telegram_api_test.go b/services/connectors/telegram_api_test.go index 9ddd8d2..d81d16d 100644 --- a/services/connectors/telegram_api_test.go +++ b/services/connectors/telegram_api_test.go @@ -10,12 +10,12 @@ import ( "testing" ) -func TestTelegramAPISendsRichMessageDraft(t *testing.T) { +func TestTelegramAPISendsMessageDraft(t *testing.T) { t.Parallel() const token = "123456:test-token" server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if got, want := r.URL.Path, "/bot"+token+"/sendRichMessageDraft"; got != want { + if got, want := r.URL.Path, "/bot"+token+"/sendMessageDraft"; got != want { t.Errorf("path = %q, want %q", got, want) } if got, want := r.Header.Get("Content-Type"), "application/json"; got != want { @@ -29,9 +29,7 @@ func TestTelegramAPISendsRichMessageDraft(t *testing.T) { want := map[string]any{ "chat_id": float64(42), "draft_id": float64(77), - "rich_message": map[string]any{ - "markdown": "**working**", - }, + "text": "working", } if !equalJSON(payload, want) { t.Errorf("payload = %#v, want %#v", payload, want) @@ -42,15 +40,13 @@ func TestTelegramAPISendsRichMessageDraft(t *testing.T) { defer server.Close() api := newTelegramHTTPAPI(token, server.Client(), server.URL) - err := api.sendRichMessageDraft(context.Background(), telegramRichMessageDraft{ + err := api.sendMessageDraft(context.Background(), telegramMessageDraft{ ChatID: 42, DraftID: 77, - RichMessage: telegramInputRichMessage{ - Markdown: "**working**", - }, + Text: "working", }) if err != nil { - t.Fatalf("sendRichMessageDraft() error = %v", err) + t.Fatalf("sendMessageDraft() error = %v", err) } } @@ -125,16 +121,16 @@ func TestTelegramAPIReturnsRetryAfterAndRedactsEchoedToken(t *testing.T) { defer server.Close() api := newTelegramHTTPAPI(token, server.Client(), server.URL) - err := api.sendRichMessageDraft(context.Background(), telegramRichMessageDraft{ - ChatID: 1, DraftID: 9, RichMessage: telegramInputRichMessage{Markdown: "text"}, + err := api.sendMessageDraft(context.Background(), telegramMessageDraft{ + ChatID: 1, DraftID: 9, Text: "text", }) if err == nil { - t.Fatal("sendRichMessageDraft() error = nil, want Telegram API error") + t.Fatal("sendMessageDraft() error = nil, want Telegram API error") } if strings.Contains(err.Error(), token) { t.Fatalf("error leaked bot token: %q", err) } - if got := err.Error(); !strings.Contains(got, "sendRichMessageDraft") || !strings.Contains(got, "send failed for [REDACTED] then [REDACTED]") { + if got := err.Error(); !strings.Contains(got, "sendMessageDraft") || !strings.Contains(got, "send failed for [REDACTED] then [REDACTED]") { t.Errorf("error = %q, want method and redacted description", got) } @@ -151,11 +147,11 @@ func TestTelegramAPIRejectsZeroDraftID(t *testing.T) { t.Parallel() api := newTelegramHTTPAPI("token", http.DefaultClient, "http://unused.invalid") - err := api.sendRichMessageDraft(context.Background(), telegramRichMessageDraft{ - ChatID: 1, RichMessage: telegramInputRichMessage{Markdown: "text"}, + err := api.sendMessageDraft(context.Background(), telegramMessageDraft{ + ChatID: 1, Text: "text", }) if err == nil || !strings.Contains(err.Error(), "draft_id must be nonzero") { - t.Fatalf("sendRichMessageDraft() error = %v, want nonzero draft ID error", err) + t.Fatalf("sendMessageDraft() error = %v, want nonzero draft ID error", err) } } diff --git a/services/connectors/telegram_integration_test.go b/services/connectors/telegram_integration_test.go index 81881b7..aead661 100644 --- a/services/connectors/telegram_integration_test.go +++ b/services/connectors/telegram_integration_test.go @@ -143,15 +143,14 @@ func TestTelegramStreamingJobCancellationStillAllowsFinalDelivery(t *testing.T) } } -func TestTelegramDeliveryPreviewCancellationAbortsInflightLegacyEdit(t *testing.T) { +func TestTelegramTurnCancellationAbortsInflightLegacyEdit(t *testing.T) { tg := &Telegram{placeholders: map[string]int{"job": 10}} b := &contextBlockingTelegramBot{ started: make(chan struct{}), returned: make(chan struct{}), } delivery := tg.telegramDelivery(context.Background(), b, -1, 7, "job", 10) - job, session := telegramNewJobWithStream(t.Context(), &telegramStreamAPI{}, -1, false, delivery, nil, "job", nil) - defer session.Close() + _, session := telegramNewJobWithStream(t.Context(), &telegramStreamAPI{}, -1, false, delivery, nil, "job", nil) session.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "working"}) select { @@ -160,7 +159,7 @@ func TestTelegramDeliveryPreviewCancellationAbortsInflightLegacyEdit(t *testing. t.Fatal("legacy preview edit did not start") } - job.Cancel() + session.Close() select { case <-b.returned: case <-time.After(time.Second): @@ -174,7 +173,7 @@ type cancellingTelegramAPI struct { calls atomic.Int32 } -func (a *cancellingTelegramAPI) sendRichMessageDraft(ctx context.Context, _ telegramRichMessageDraft) error { +func (a *cancellingTelegramAPI) sendMessageDraft(ctx context.Context, _ telegramMessageDraft) error { if a.calls.Add(1) == 1 { close(a.started) } @@ -184,16 +183,15 @@ func (a *cancellingTelegramAPI) sendRichMessageDraft(ctx context.Context, _ tele } func (*cancellingTelegramAPI) sendRichMessage(context.Context, telegramRichMessage) error { return nil } -func TestTelegramTrackedJobCancellationAbortsInflightDraftRequest(t *testing.T) { +func TestTelegramTurnCancellationAbortsInflightDraftRequest(t *testing.T) { api := &cancellingTelegramAPI{started: make(chan struct{}), returned: make(chan struct{})} - job, session := telegramNewJobWithStream(t.Context(), api, 1, true, telegramStreamDelivery{}, nil, "job", nil) - defer session.Close() + _, session := telegramNewJobWithStream(t.Context(), api, 1, true, telegramStreamDelivery{}, nil, "job", nil) select { case <-api.started: case <-time.After(time.Second): t.Fatal("draft request did not start") } - job.Cancel() + session.Close() select { case <-api.returned: case <-time.After(time.Second): @@ -202,15 +200,28 @@ func TestTelegramTrackedJobCancellationAbortsInflightDraftRequest(t *testing.T) session.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "must not restart previews"}) time.Sleep(telegramStreamInterval + 50*time.Millisecond) if got := api.calls.Load(); got != 1 { - t.Fatalf("draft calls after job cancellation = %d, want 1", got) + t.Fatalf("draft calls after turn cancellation = %d, want 1", got) } - select { - case <-session.done: - t.Fatal("job cancellation stopped final-delivery orchestration") - default: +} + +func TestTelegramJobCompletionDoesNotCancelPendingTurnPreview(t *testing.T) { + api := &telegramStreamAPI{} + job, session := telegramNewJobWithStream(t.Context(), api, 1, true, telegramStreamDelivery{}, nil, "job", nil) + defer session.Close() + waitTelegramStream(t, func() bool { drafts, _ := api.snapshot(); return len(drafts) == 1 }) + + session.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "completed answer"}) + // Agent.consumeJob cancels its job context when computation completes. The + // Telegram consumer belongs to the surrounding turn and must remain alive + // long enough for the handler to drain it. + job.Cancel() + if err := session.Flush(); err != nil { + t.Fatalf("Flush after job completion: %v", err) } - if err := session.Finalize("answer after cancellation", nil); err != nil { - t.Fatalf("Finalize after cancellation: %v", err) + + drafts, _ := api.snapshot() + if got := drafts[len(drafts)-1].Text; got != "completed answer"+telegramStreamCursor { + t.Fatalf("flushed draft = %q, want completed answer", got) } } @@ -224,7 +235,7 @@ type contextRecordingTelegramAPI struct { contexts []context.Context } -func (*contextRecordingTelegramAPI) sendRichMessageDraft(context.Context, telegramRichMessageDraft) error { +func (*contextRecordingTelegramAPI) sendMessageDraft(context.Context, telegramMessageDraft) error { return nil } @@ -239,7 +250,7 @@ func (a *contextRecordingTelegramAPI) sendRichMessage(ctx context.Context, _ tel return errors.New("force fallback delivery") } -func (*orderedTelegramAPI) sendRichMessageDraft(context.Context, telegramRichMessageDraft) error { +func (*orderedTelegramAPI) sendMessageDraft(context.Context, telegramMessageDraft) error { return nil } func (a *orderedTelegramAPI) sendRichMessage(_ context.Context, m telegramRichMessage) error { @@ -356,7 +367,7 @@ func TestTelegramAskOptionsAttachMatchingSession(t *testing.T) { t.Fatal(err) } drafts, _ := api.snapshot() - if got := drafts[len(drafts)-1].RichMessage.Markdown; got != "hello" { + if got := drafts[len(drafts)-1].Text; got != "hello"+telegramStreamCursor { t.Fatalf("preview = %q, want hello", got) } } diff --git a/services/connectors/telegram_stream.go b/services/connectors/telegram_stream.go index ac798c1..a0d2aab 100644 --- a/services/connectors/telegram_stream.go +++ b/services/connectors/telegram_stream.go @@ -8,10 +8,12 @@ import ( "time" "github.com/mudler/cogito" + "github.com/mudler/xlog" ) const telegramStreamInterval = 400 * time.Millisecond const telegramDraftHeartbeatInterval = 25 * time.Second +const telegramStreamCursor = " ▉" type telegramStreamDelivery struct { editPreview func(context.Context, int64, string) error @@ -271,7 +273,10 @@ func (s *telegramStreamSession) run() { schedule(when) continue } - attempted, retry, _ := s.deliverPreview() + attempted, retry, err := s.deliverPreview() + if err != nil && retry == 0 && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + xlog.Warn("Telegram preview delivery failed", "chat_id", s.chatID, "error", err) + } if attempted { nextPreview = time.Now().Add(telegramStreamInterval) } @@ -298,7 +303,7 @@ func (s *telegramStreamSession) previewSnapshot() (string, uint64, bool, bool) { return "", 0, false, false } if s.thinkingPending { - return telegramThinkingMessage, s.version, true, true + return telegramThinkingMessage + telegramStreamCursor, s.version, true, true } text := s.content if text == "" { @@ -307,7 +312,8 @@ func (s *telegramStreamSession) previewSnapshot() (string, uint64, bool, bool) { text = telegramThinkingMessage } } - return telegramPreviewTail(text, telegramMaxMessageLength), s.version, true, false + limit := telegramMaxMessageLength - len([]rune(telegramStreamCursor)) + return telegramPreviewTail(text, limit) + telegramStreamCursor, s.version, true, false } func telegramPreviewTail(text string, limit int) string { @@ -341,12 +347,16 @@ func (s *telegramStreamSession) deliverPreview() (bool, time.Duration, error) { } var err error if s.private { - err = s.api.sendRichMessageDraft(s.ctx, telegramRichMessageDraft{ChatID: s.chatID, DraftID: s.draftID, RichMessage: telegramInputRichMessage{Markdown: text}}) + err = s.api.sendMessageDraft(s.ctx, telegramMessageDraft{ChatID: s.chatID, DraftID: s.draftID, Text: text}) var apiErr *telegramAPIError if errors.As(err, &apiErr) && apiErr.RetryAfter > 0 { return true, time.Duration(apiErr.RetryAfter) * time.Second, err } if err != nil { + if s.ctx.Err() != nil { + return true, 0, err + } + xlog.Warn("Telegram sendMessageDraft failed; falling back to editable message", "chat_id", s.chatID, "error", err) s.private = false if s.delivery.editPreview != nil { err = s.delivery.editPreview(s.ctx, s.chatID, text) diff --git a/services/connectors/telegram_stream_test.go b/services/connectors/telegram_stream_test.go index 7a6fe63..1b643c6 100644 --- a/services/connectors/telegram_stream_test.go +++ b/services/connectors/telegram_stream_test.go @@ -15,7 +15,7 @@ import ( type telegramStreamAPI struct { mu sync.Mutex - drafts []telegramRichMessageDraft + drafts []telegramMessageDraft finals []telegramRichMessage draftErr func(int) error finalErr func(int) error @@ -24,7 +24,7 @@ type telegramStreamAPI struct { block time.Duration } -func (a *telegramStreamAPI) sendRichMessageDraft(_ context.Context, draft telegramRichMessageDraft) error { +func (a *telegramStreamAPI) sendMessageDraft(_ context.Context, draft telegramMessageDraft) error { n := a.inCall.Add(1) defer a.inCall.Add(-1) for old := a.maxCalls.Load(); n > old && !a.maxCalls.CompareAndSwap(old, n); old = a.maxCalls.Load() { @@ -57,10 +57,10 @@ func (a *telegramStreamAPI) sendRichMessage(_ context.Context, final telegramRic return nil } -func (a *telegramStreamAPI) snapshot() ([]telegramRichMessageDraft, []telegramRichMessage) { +func (a *telegramStreamAPI) snapshot() ([]telegramMessageDraft, []telegramRichMessage) { a.mu.Lock() defer a.mu.Unlock() - return append([]telegramRichMessageDraft(nil), a.drafts...), append([]telegramRichMessage(nil), a.finals...) + return append([]telegramMessageDraft(nil), a.drafts...), append([]telegramRichMessage(nil), a.finals...) } func waitTelegramStream(t *testing.T, condition func() bool) { @@ -93,7 +93,7 @@ func TestTelegramStreamPrivateUsesStableDraftAndRateLimitsSerializedCalls(t *tes if drafts[0].DraftID == 0 || drafts[1].DraftID != drafts[0].DraftID { t.Fatalf("draft IDs = %d, %d, want same nonzero ID", drafts[0].DraftID, drafts[1].DraftID) } - if drafts[0].RichMessage.Markdown != telegramThinkingMessage || drafts[1].RichMessage.Markdown != "one two three" { + if drafts[0].Text != telegramThinkingMessage+telegramStreamCursor || drafts[1].Text != "one two three"+telegramStreamCursor { t.Fatalf("drafts = %#v", drafts) } if api.maxCalls.Load() != 1 { @@ -112,10 +112,10 @@ func TestTelegramStreamAlwaysDeliversThinkingBeforeImmediateContent(t *testing.T waitTelegramStream(t, func() bool { drafts, _ := api.snapshot(); return len(drafts) == 2 }) drafts, _ := api.snapshot() - if got := drafts[0].RichMessage.Markdown; got != telegramThinkingMessage { + if got := drafts[0].Text; got != telegramThinkingMessage+telegramStreamCursor { t.Fatalf("initial draft = %q, want thinking draft", got) } - if got := drafts[1].RichMessage.Markdown; got != "immediate" { + if got := drafts[1].Text; got != "immediate"+telegramStreamCursor { t.Fatalf("content draft = %q, want immediate content without another event", got) } } @@ -123,7 +123,7 @@ func TestTelegramStreamAlwaysDeliversThinkingBeforeImmediateContent(t *testing.T func TestTelegramStreamRetryAfterRetainsLatestPreview(t *testing.T) { api := &telegramStreamAPI{draftErr: func(i int) error { if i == 2 { - return &telegramAPIError{Method: "sendRichMessageDraft", ErrorCode: 429, RetryAfter: 1} + return &telegramAPIError{Method: "sendMessageDraft", ErrorCode: 429, RetryAfter: 1} } return nil }} @@ -135,7 +135,7 @@ func TestTelegramStreamRetryAfterRetainsLatestPreview(t *testing.T) { s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: " latest"}) waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 3 }) drafts, _ := api.snapshot() - if got := drafts[2].RichMessage.Markdown; got != "first latest" { + if got := drafts[2].Text; got != "first latest"+telegramStreamCursor { t.Fatalf("retried preview = %q", got) } } @@ -155,7 +155,7 @@ func TestTelegramStreamEditRetryAfterReschedulesPendingPreviewWithoutNewContent( select { case text := <-got: - if text != telegramThinkingMessage { + if text != telegramThinkingMessage+telegramStreamCursor { t.Fatalf("retried preview = %q, want thinking preview", text) } case <-time.After(2 * time.Second): @@ -187,7 +187,7 @@ func TestTelegramStreamNativeFailureFallsBackOnlyForThatSession(t *testing.T) { waitTelegramStream(t, func() bool { mu.Lock() defer mu.Unlock() - return len(edits[1]) > 0 && edits[1][len(edits[1])-1] == "fallback" + return len(edits[1]) > 0 && edits[1][len(edits[1])-1] == "fallback"+telegramStreamCursor }) waitTelegramStream(t, func() bool { d, _ := healthy.snapshot(); return len(d) >= 2 }) mu.Lock() @@ -237,7 +237,7 @@ func TestTelegramStreamFlushDeliversPendingContentBeforeReturning(t *testing.T) if len(drafts) != 2 { t.Fatalf("drafts when Flush returned = %d, want pending content delivered", len(drafts)) } - if got := drafts[1].RichMessage.Markdown; got != "pending answer" { + if got := drafts[1].Text; got != "pending answer"+telegramStreamCursor { t.Fatalf("flushed preview = %q, want pending answer", got) } } @@ -319,8 +319,8 @@ func TestTelegramStreamPreviewUsesUTF8SafeTail(t *testing.T) { t.Fatal(err) } drafts, _ := api.snapshot() - got := drafts[len(drafts)-1].RichMessage.Markdown - if len([]rune(got)) != telegramMaxMessageLength || got != strings.Repeat("🙂", telegramMaxMessageLength) { + got := drafts[len(drafts)-1].Text + if len([]rune(got)) != telegramMaxMessageLength || got != strings.Repeat("🙂", telegramMaxMessageLength-len([]rune(telegramStreamCursor)))+telegramStreamCursor { t.Fatalf("preview rune length = %d, want tail of %d", len([]rune(got)), telegramMaxMessageLength) } if s.content != full { @@ -338,7 +338,7 @@ func TestTelegramStreamPrivateShowsPublishedStatusBeforeContent(t *testing.T) { t.Fatal(err) } drafts, _ := api.snapshot() - if got := drafts[len(drafts)-1].RichMessage.Markdown; got != "checking sources" { + if got := drafts[len(drafts)-1].Text; got != "checking sources"+telegramStreamCursor { t.Fatalf("status = %q", got) } s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "answer"}) @@ -346,7 +346,7 @@ func TestTelegramStreamPrivateShowsPublishedStatusBeforeContent(t *testing.T) { t.Fatal(err) } drafts, _ = api.snapshot() - if got := drafts[len(drafts)-1].RichMessage.Markdown; got != "answer" { + if got := drafts[len(drafts)-1].Text; got != "answer"+telegramStreamCursor { t.Fatalf("answer = %q", got) } } @@ -402,7 +402,7 @@ func TestTelegramStreamGroupEditsPlaceholder(t *testing.T) { defer s.Close() select { case text := <-got: - if text != telegramThinkingMessage { + if text != telegramThinkingMessage+telegramStreamCursor { t.Fatalf("initial edit = %q", text) } case <-time.After(time.Second): @@ -411,7 +411,7 @@ func TestTelegramStreamGroupEditsPlaceholder(t *testing.T) { s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "group answer"}) select { case text := <-got: - if text != "group answer" { + if text != "group answer"+telegramStreamCursor { t.Fatalf("content edit = %q", text) } case <-time.After(time.Second): @@ -483,7 +483,7 @@ func TestTelegramStreamNativeDraftHeartbeatAndClose(t *testing.T) { s := newTelegramStreamSessionWithHeartbeat(context.Background(), api, 42, true, telegramStreamDelivery{}, 20*time.Millisecond) waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) >= 2 }) drafts, _ := api.snapshot() - if drafts[0].DraftID != drafts[1].DraftID || drafts[0].RichMessage.Markdown != drafts[1].RichMessage.Markdown { + if drafts[0].DraftID != drafts[1].DraftID || drafts[0].Text != drafts[1].Text { t.Fatalf("heartbeats changed draft: %#v", drafts[:2]) } s.Close()