diff --git a/src/telegram/bot-message-dispatch.test.ts b/src/telegram/bot-message-dispatch.test.ts index 5f59feb77a..882d80997a 100644 --- a/src/telegram/bot-message-dispatch.test.ts +++ b/src/telegram/bot-message-dispatch.test.ts @@ -583,6 +583,70 @@ describe("dispatchTelegramMessage draft streaming", () => { ); }); + it("does not duplicate reasoning final after reasoning end in partial mode", async () => { + let reasoningMessageId: number | undefined = 111; + const reasoningDraftStream = { + update: vi.fn(), + flush: vi.fn().mockResolvedValue(undefined), + messageId: vi.fn().mockImplementation(() => reasoningMessageId), + clear: vi.fn().mockResolvedValue(undefined), + stop: vi.fn().mockResolvedValue(undefined), + forceNewMessage: vi.fn().mockImplementation(() => { + reasoningMessageId = undefined; + }), + }; + const answerDraftStream = createDraftStream(999); + createTelegramDraftStream + .mockImplementationOnce(() => answerDraftStream) + .mockImplementationOnce(() => reasoningDraftStream); + dispatchReplyWithBufferedBlockDispatcher.mockImplementation( + async ({ dispatcherOptions, replyOptions }) => { + await replyOptions?.onReasoningStream?.({ text: "Reasoning:\n_step one_" }); + await replyOptions?.onReasoningEnd?.(); + await dispatcherOptions.deliver( + { text: "Reasoning:\n_step one expanded_" }, + { kind: "final" }, + ); + return { queuedFinal: true }; + }, + ); + deliverReplies.mockResolvedValue({ delivered: true }); + editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "111" }); + + await dispatchWithContext({ context: createContext(), streamMode: "partial" }); + + expect(editMessageTelegram).toHaveBeenCalledWith( + 123, + 111, + "Reasoning:\n_step one expanded_", + expect.any(Object), + ); + expect(deliverReplies).not.toHaveBeenCalled(); + }); + + it("splits reasoning preview only when next reasoning block starts in partial mode", async () => { + const { reasoningDraftStream } = setupDraftStreams({ + answerMessageId: 999, + reasoningMessageId: 111, + }); + dispatchReplyWithBufferedBlockDispatcher.mockImplementation( + async ({ dispatcherOptions, replyOptions }) => { + await replyOptions?.onReasoningStream?.({ text: "Reasoning:\n_first block_" }); + await replyOptions?.onReasoningEnd?.(); + await replyOptions?.onPartialReply?.({ text: "checking files..." }); + await replyOptions?.onReasoningStream?.({ text: "Reasoning:\n_second block_" }); + await dispatcherOptions.deliver({ text: "Done" }, { kind: "final" }); + return { queuedFinal: true }; + }, + ); + deliverReplies.mockResolvedValue({ delivered: true }); + editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "999" }); + + await dispatchWithContext({ context: createContext(), streamMode: "partial" }); + + expect(reasoningDraftStream.forceNewMessage).toHaveBeenCalledTimes(1); + }); + it("does not edit preview message when final payload is an error", async () => { const draftStream = createDraftStream(999); createTelegramDraftStream.mockReturnValue(draftStream); diff --git a/src/telegram/bot-message-dispatch.ts b/src/telegram/bot-message-dispatch.ts index 06054da705..f7f9a6bcb7 100644 --- a/src/telegram/bot-message-dispatch.ts +++ b/src/telegram/bot-message-dispatch.ts @@ -169,6 +169,7 @@ export const dispatchTelegramMessage = async ({ hasStreamedMessage: false, chunker: reasoningDraftChunker, }; + let splitReasoningOnNextStream = false; const resetDraftLaneState = (lane: DraftLaneState) => { lane.lastPartialText = ""; lane.draftText = ""; @@ -485,7 +486,17 @@ export const dispatchTelegramMessage = async ({ ? (payload) => updateDraftFromPartial(answerLane, payload.text) : undefined, onReasoningStream: reasoningLane.stream - ? (payload) => updateDraftFromPartial(reasoningLane, payload.text) + ? (payload) => { + // In partial mode, split between reasoning blocks only when the + // next reasoning stream starts. Splitting at reasoning-end can + // orphan the active preview and cause duplicate reasoning sends. + if (streamMode === "partial" && splitReasoningOnNextStream) { + reasoningLane.stream?.forceNewMessage(); + resetDraftLaneState(reasoningLane); + splitReasoningOnNextStream = false; + } + updateDraftFromPartial(reasoningLane, payload.text); + } : undefined, onAssistantMessageStart: answerLane.stream ? () => { @@ -498,12 +509,16 @@ export const dispatchTelegramMessage = async ({ : undefined, onReasoningEnd: reasoningLane.stream ? () => { - // Reasoning should appear in its own lane and each block should keep - // a separate bubble for readability. - if (reasoningLane.hasStreamedMessage) { - reasoningLane.stream?.forceNewMessage(); + // Block mode keeps hard message boundaries at reasoning-end. + if (streamMode === "block") { + if (reasoningLane.hasStreamedMessage) { + reasoningLane.stream?.forceNewMessage(); + } + resetDraftLaneState(reasoningLane); + return; } - resetDraftLaneState(reasoningLane); + // Partial mode splits when/if a later reasoning block begins. + splitReasoningOnNextStream = reasoningLane.hasStreamedMessage; } : undefined, onModelSelected,