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 <mudler@localai.io>

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
 }

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 {

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)
 	}
 }
 

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)
 	}
 }

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)

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()