fix: stabilize Telegram draft boundaries and suppress NO_REPLY lead leaks (#33169)

* fix: stabilize telegram draft stream message boundaries

* fix: suppress NO_REPLY lead-fragment leaks

* fix: keep underscore guard for non-NO_REPLY prefixes

* fix: skip assistant-start rotation only after real lane rotation

* fix: preserve finalized state when pre-rotation does not force

* fix: reset finalized preview state on message-start boundary

* fix: document Telegram draft boundary + NO_REPLY reliability updates (#33169) (thanks @obviyus)
This commit is contained in:
Ayaan Zaidi
2026-03-03 22:49:33 +05:30
committed by GitHub
parent a7a9a3d3c8
commit 3d998828b9
8 changed files with 212 additions and 137 deletions

View File

@@ -444,6 +444,133 @@ describe("dispatchTelegramMessage draft streaming", () => {
expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
});
it("rotates before a late second-message partial so finalized preview is not overwritten", async () => {
const answerDraftStream = createSequencedDraftStream(1001);
const reasoningDraftStream = createDraftStream();
createTelegramDraftStream
.mockImplementationOnce(() => answerDraftStream)
.mockImplementationOnce(() => reasoningDraftStream);
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
async ({ dispatcherOptions, replyOptions }) => {
await replyOptions?.onPartialReply?.({ text: "Message A partial" });
await dispatcherOptions.deliver({ text: "Message A final" }, { kind: "final" });
// Simulate provider ordering bug: first chunk arrives before message-start callback.
await replyOptions?.onPartialReply?.({ text: "Message B early" });
await replyOptions?.onAssistantMessageStart?.();
await replyOptions?.onPartialReply?.({ text: "Message B partial" });
await dispatcherOptions.deliver({ text: "Message B final" }, { kind: "final" });
return { queuedFinal: true };
},
);
deliverReplies.mockResolvedValue({ delivered: true });
editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "1001" });
await dispatchWithContext({ context: createContext(), streamMode: "partial" });
expect(answerDraftStream.forceNewMessage).toHaveBeenCalledTimes(1);
expect(answerDraftStream.update).toHaveBeenNthCalledWith(2, "Message B early");
const boundaryRotationOrder = answerDraftStream.forceNewMessage.mock.invocationCallOrder[0];
const secondUpdateOrder = answerDraftStream.update.mock.invocationCallOrder[1];
expect(boundaryRotationOrder).toBeLessThan(secondUpdateOrder);
expect(editMessageTelegram).toHaveBeenNthCalledWith(
1,
123,
1001,
"Message A final",
expect.any(Object),
);
expect(editMessageTelegram).toHaveBeenNthCalledWith(
2,
123,
1002,
"Message B final",
expect.any(Object),
);
});
it("does not skip message-start rotation when pre-rotation did not force a new message", async () => {
const answerDraftStream = createSequencedDraftStream(1002);
answerDraftStream.setMessageId(1001);
const reasoningDraftStream = createDraftStream();
createTelegramDraftStream
.mockImplementationOnce(() => answerDraftStream)
.mockImplementationOnce(() => reasoningDraftStream);
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
async ({ dispatcherOptions, replyOptions }) => {
// First message has only final text (no streamed partials), so answer lane
// reaches finalized state with hasStreamedMessage still false.
await dispatcherOptions.deliver({ text: "Message A final" }, { kind: "final" });
// Provider ordering bug: next message partial arrives before message-start.
await replyOptions?.onPartialReply?.({ text: "Message B early" });
await replyOptions?.onAssistantMessageStart?.();
await replyOptions?.onPartialReply?.({ text: "Message B partial" });
await dispatcherOptions.deliver({ text: "Message B final" }, { kind: "final" });
return { queuedFinal: true };
},
);
deliverReplies.mockResolvedValue({ delivered: true });
editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "1001" });
const bot = createBot();
await dispatchWithContext({ context: createContext(), streamMode: "partial", bot });
// Early pre-rotation could not force (no streamed partials yet), so the
// real assistant message_start must still rotate once.
expect(answerDraftStream.forceNewMessage).toHaveBeenCalledTimes(1);
expect(answerDraftStream.update).toHaveBeenNthCalledWith(1, "Message B early");
expect(answerDraftStream.update).toHaveBeenNthCalledWith(2, "Message B partial");
const earlyUpdateOrder = answerDraftStream.update.mock.invocationCallOrder[0];
const boundaryRotationOrder = answerDraftStream.forceNewMessage.mock.invocationCallOrder[0];
const secondUpdateOrder = answerDraftStream.update.mock.invocationCallOrder[1];
expect(earlyUpdateOrder).toBeLessThan(boundaryRotationOrder);
expect(boundaryRotationOrder).toBeLessThan(secondUpdateOrder);
expect(editMessageTelegram).toHaveBeenNthCalledWith(
1,
123,
1001,
"Message A final",
expect.any(Object),
);
expect(editMessageTelegram).toHaveBeenNthCalledWith(
2,
123,
1002,
"Message B final",
expect.any(Object),
);
expect((bot.api.deleteMessage as ReturnType<typeof vi.fn>).mock.calls).toHaveLength(0);
});
it("does not trigger late pre-rotation mid-message after an explicit assistant message start", async () => {
const answerDraftStream = createDraftStream(1001);
const reasoningDraftStream = createDraftStream();
createTelegramDraftStream
.mockImplementationOnce(() => answerDraftStream)
.mockImplementationOnce(() => reasoningDraftStream);
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
async ({ dispatcherOptions, replyOptions }) => {
// Message A finalizes without streamed partials.
await dispatcherOptions.deliver({ text: "Message A final" }, { kind: "final" });
// Message B starts normally before partials.
await replyOptions?.onAssistantMessageStart?.();
await replyOptions?.onPartialReply?.({ text: "Message B first chunk" });
await replyOptions?.onPartialReply?.({ text: "Message B second chunk" });
await dispatcherOptions.deliver({ text: "Message B final" }, { kind: "final" });
return { queuedFinal: true };
},
);
deliverReplies.mockResolvedValue({ delivered: true });
editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "1001" });
await dispatchWithContext({ context: createContext(), streamMode: "partial" });
// The explicit message_start boundary must clear finalized state so
// same-message partials do not force a new preview mid-stream.
expect(answerDraftStream.forceNewMessage).not.toHaveBeenCalled();
expect(answerDraftStream.update).toHaveBeenNthCalledWith(1, "Message B first chunk");
expect(answerDraftStream.update).toHaveBeenNthCalledWith(2, "Message B second chunk");
});
it("finalizes multi-message assistant stream to matching preview messages in order", async () => {
const answerDraftStream = createSequencedDraftStream(1001);
const reasoningDraftStream = createDraftStream();