Skip to content

Commit 0cab701

Browse files
obviyusobviyus
committed
fix(telegram): use partial stream deltas
1 parent 8151d4d commit 0cab701

6 files changed

Lines changed: 106 additions & 67 deletions

File tree

‎extensions/telegram/src/bot-message-dispatch.test.ts‎

Lines changed: 43 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -984,14 +984,21 @@ describe("dispatchTelegramMessage draft streaming", () => {
984984
expect(editMessageTelegram).not.toHaveBeenCalled();
985985
});
986986

987-
it("coalesces delta-shaped partial fragments while preserving the first-preview debounce", async () => {
987+
it("applies partial deltas while preserving the first-preview debounce", async () => {
988988
const { answerDraftStream } = setupDraftStreams({ answerMessageId: 2001 });
989989
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
990990
async ({ dispatcherOptions, replyOptions }) => {
991-
await replyOptions?.onPartialReply?.({ text: "Streaming " });
992-
await replyOptions?.onPartialReply?.({ text: "previews " });
993991
await replyOptions?.onPartialReply?.({
994-
text: "are useful because they show progress.",
992+
text: "Streaming ",
993+
delta: "Streaming ",
994+
});
995+
await replyOptions?.onPartialReply?.({
996+
text: "Streaming previews ",
997+
delta: "previews ",
998+
});
999+
await replyOptions?.onPartialReply?.({
1000+
text: "Streaming previews are useful because they show progress.",
1001+
delta: "are useful because they show progress.",
9951002
});
9961003
await dispatcherOptions.deliver(
9971004
{ text: "Streaming previews are useful because they show progress." },
@@ -1025,13 +1032,43 @@ describe("dispatchTelegramMessage draft streaming", () => {
10251032
expect(deliverReplies).not.toHaveBeenCalled();
10261033
});
10271034

1035+
it("replaces non-prefix partial snapshots instead of appending them", async () => {
1036+
const { answerDraftStream } = setupDraftStreams({ answerMessageId: 2001 });
1037+
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
1038+
async ({ dispatcherOptions, replyOptions }) => {
1039+
await replyOptions?.onPartialReply?.({
1040+
text: "Working...",
1041+
delta: "Working...",
1042+
});
1043+
await replyOptions?.onPartialReply?.({
1044+
text: "Done.",
1045+
delta: "",
1046+
replace: true,
1047+
});
1048+
await dispatcherOptions.deliver({ text: "Done." }, { kind: "final" });
1049+
return { queuedFinal: true };
1050+
},
1051+
);
1052+
1053+
await dispatchWithContext({
1054+
context: createContext(),
1055+
streamMode: "partial",
1056+
telegramCfg: { streaming: { mode: "partial" } },
1057+
});
1058+
1059+
expect(answerDraftStream.update).toHaveBeenNthCalledWith(1, "Working...");
1060+
expect(answerDraftStream.update).toHaveBeenNthCalledWith(2, "Done.");
1061+
expect(answerDraftStream.update).toHaveBeenLastCalledWith("Done.");
1062+
expect(deliverReplies).not.toHaveBeenCalled();
1063+
});
1064+
10281065
it("does not coalesce answer partial fragments with tool progress drafts", async () => {
10291066
const { answerDraftStream } = setupDraftStreams({ answerMessageId: 2001 });
10301067
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
10311068
async ({ dispatcherOptions, replyOptions }) => {
10321069
await replyOptions?.onToolStart?.({ name: "exec", phase: "start" });
1033-
await replyOptions?.onPartialReply?.({ text: "Done " });
1034-
await replyOptions?.onPartialReply?.({ text: "answer" });
1070+
await replyOptions?.onPartialReply?.({ text: "Done ", delta: "Done " });
1071+
await replyOptions?.onPartialReply?.({ text: "Done answer", delta: "answer" });
10351072
await dispatcherOptions.deliver({ text: "Done answer." }, { kind: "final" });
10361073
return { queuedFinal: true };
10371074
},

‎extensions/telegram/src/bot-message-dispatch.ts‎

Lines changed: 49 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -110,49 +110,22 @@ const silentReplyDispatchLogger = createSubsystemLogger("telegram/silent-reply-d
110110
/** Minimum chars before sending first streaming message (improves push notification UX) */
111111
const DRAFT_MIN_INITIAL_CHARS = 30;
112112

113-
function appendWithOverlap(previous: string, fragment: string): string {
114-
const maxOverlap = Math.min(previous.length, fragment.length);
115-
for (let overlap = maxOverlap; overlap > 0; overlap -= 1) {
116-
if (previous.endsWith(fragment.slice(0, overlap))) {
117-
return `${previous}${fragment.slice(overlap)}`;
118-
}
119-
}
120-
return `${previous}${fragment}`;
121-
}
122-
123-
function looksLikeDraftDeltaFragment(previous: string, text: string): boolean {
124-
if (!previous || !text) {
125-
return false;
126-
}
127-
if (text.startsWith(previous) || previous.startsWith(text)) {
128-
return false;
129-
}
130-
if (/^\s/.test(text)) {
131-
return true;
132-
}
133-
if (previous.length < DRAFT_MIN_INITIAL_CHARS) {
134-
return true;
135-
}
136-
if (/\s$/.test(previous) && text.length <= DRAFT_MIN_INITIAL_CHARS) {
137-
return true;
138-
}
139-
return text.length <= Math.max(16, Math.floor(previous.length / 2));
140-
}
113+
type DraftPartialTextUpdate = {
114+
text: string;
115+
delta?: string;
116+
replace?: true;
117+
};
141118

142-
function resolveDraftPartialText(previous: string, text: string): string | undefined {
143-
if (!previous) {
144-
return text;
145-
}
146-
if (text === previous) {
147-
return undefined;
148-
}
149-
if (text.startsWith(previous)) {
150-
return text;
151-
}
152-
if (previous.startsWith(text) && text.length < previous.length) {
119+
function resolveDraftPartialText(
120+
previous: string,
121+
update: DraftPartialTextUpdate,
122+
): string | undefined {
123+
const nextText =
124+
update.replace || update.delta === undefined ? update.text : `${previous}${update.delta}`;
125+
if (nextText === previous) {
153126
return undefined;
154127
}
155-
return looksLikeDraftDeltaFragment(previous, text) ? appendWithOverlap(previous, text) : text;
128+
return nextText;
156129
}
157130

158131
async function resolveStickerVisionSupport(cfg: OpenClawConfig, agentId: string) {
@@ -714,23 +687,36 @@ export const dispatchTelegramMessage = async ({
714687
});
715688
return draftLaneEventQueue;
716689
};
717-
type SplitLaneSegment = { lane: LaneName; text: string };
690+
type SplitLaneSegment = { lane: LaneName; update: DraftPartialTextUpdate };
718691
type SplitLaneSegmentsResult = {
719692
segments: SplitLaneSegment[];
720693
suppressedReasoningOnly: boolean;
721694
};
722695
const splitTextIntoLaneSegments = (
723-
text?: string,
696+
update: { text?: string; delta?: string; replace?: true },
724697
isReasoning?: boolean,
725698
): SplitLaneSegmentsResult => {
726-
const split = splitTelegramReasoningText(text, isReasoning);
699+
const split = splitTelegramReasoningText(update.text, isReasoning);
700+
const splitSegments: Array<{ lane: LaneName; text: string }> = [];
701+
const useDelta = !update.replace && update.delta !== undefined;
727702
const segments: SplitLaneSegment[] = [];
728703
const suppressReasoning = resolvedReasoningLevel === "off";
729704
if (split.reasoningText && !suppressReasoning) {
730-
segments.push({ lane: "reasoning", text: split.reasoningText });
705+
splitSegments.push({ lane: "reasoning", text: split.reasoningText });
731706
}
732707
if (split.answerText) {
733-
segments.push({ lane: "answer", text: split.answerText });
708+
splitSegments.push({ lane: "answer", text: split.answerText });
709+
}
710+
for (const segment of splitSegments) {
711+
const canApplyDelta = useDelta && splitSegments.length === 1;
712+
segments.push({
713+
lane: segment.lane,
714+
update: {
715+
text: segment.text,
716+
...(canApplyDelta ? { delta: update.delta } : {}),
717+
...(update.replace ? { replace: true } : {}),
718+
},
719+
});
734720
}
735721
return {
736722
segments,
@@ -761,13 +747,13 @@ export const dispatchTelegramMessage = async ({
761747
}
762748
await rotateLaneForNewMessage(answerLane);
763749
};
764-
const updateDraftFromPartial = (lane: DraftLaneState, text: string | undefined) => {
750+
const updateDraftFromPartial = (lane: DraftLaneState, update: DraftPartialTextUpdate) => {
765751
const laneStream = lane.stream;
766-
if (!laneStream || !text) {
752+
if (!laneStream || !update.text) {
767753
return;
768754
}
769755
const previousText = lane === answerLane ? lastAnswerPartialText : lane.lastPartialText;
770-
const nextText = resolveDraftPartialText(previousText, text);
756+
const nextText = resolveDraftPartialText(previousText, update);
771757
if (!nextText) {
772758
return;
773759
}
@@ -786,8 +772,11 @@ export const dispatchTelegramMessage = async ({
786772
lane.lastPartialText = nextText;
787773
laneStream.update(nextText);
788774
};
789-
const ingestDraftLaneSegments = async (text: string | undefined, isReasoning?: boolean) => {
790-
const split = splitTextIntoLaneSegments(text, isReasoning);
775+
const ingestDraftLaneSegments = async (
776+
update: { text?: string; delta?: string; replace?: true },
777+
isReasoning?: boolean,
778+
) => {
779+
const split = splitTextIntoLaneSegments(update, isReasoning);
791780
for (const segment of split.segments) {
792781
if (segment.lane === "answer") {
793782
await prepareAnswerLaneForText();
@@ -796,7 +785,7 @@ export const dispatchTelegramMessage = async ({
796785
reasoningStepState.noteReasoningHint();
797786
reasoningStepState.noteReasoningDelivered();
798787
}
799-
updateDraftFromPartial(lanes[segment.lane], segment.text);
788+
updateDraftFromPartial(lanes[segment.lane], segment.update);
800789
}
801790
};
802791
const flushDraftLane = async (lane: DraftLaneState) => {
@@ -1207,7 +1196,10 @@ export const dispatchTelegramMessage = async ({
12071196
| { buttons?: TelegramInlineButtons }
12081197
| undefined
12091198
)?.buttons;
1210-
const split = splitTextIntoLaneSegments(payload.text, payload.isReasoning);
1199+
const split = splitTextIntoLaneSegments(
1200+
{ text: payload.text },
1201+
payload.isReasoning,
1202+
);
12111203
const segments = split.segments;
12121204
const reply = resolveSendableOutboundReplyParts(payload);
12131205
const _hasMedia = reply.hasMedia;
@@ -1241,7 +1233,7 @@ export const dispatchTelegramMessage = async ({
12411233
) {
12421234
reasoningStepState.bufferFinalAnswer({
12431235
payload,
1244-
text: segment.text,
1236+
text: segment.update.text,
12451237
bufferedGeneration: replyFenceGeneration,
12461238
});
12471239
continue;
@@ -1253,10 +1245,10 @@ export const dispatchTelegramMessage = async ({
12531245
streamMode === "progress" &&
12541246
segment.lane === "answer" &&
12551247
info.kind === "final"
1256-
? await deliverProgressModeFinalAnswer(payload, segment.text)
1248+
? await deliverProgressModeFinalAnswer(payload, segment.update.text)
12571249
: await deliverLaneText({
12581250
laneName: segment.lane,
1259-
text: segment.text,
1251+
text: segment.update.text,
12601252
payload,
12611253
infoKind: info.kind,
12621254
buttons: telegramButtons,
@@ -1351,7 +1343,7 @@ export const dispatchTelegramMessage = async ({
13511343
answerLane.stream || reasoningLane.stream
13521344
? (payload) =>
13531345
enqueueDraftLaneEvent(async () => {
1354-
await ingestDraftLaneSegments(payload.text);
1346+
await ingestDraftLaneSegments(payload);
13551347
})
13561348
: undefined,
13571349
onReasoningStream: reasoningLane.stream
@@ -1362,7 +1354,7 @@ export const dispatchTelegramMessage = async ({
13621354
resetDraftLaneState(reasoningLane);
13631355
splitReasoningOnNextStream = false;
13641356
}
1365-
await ingestDraftLaneSegments(payload.text, true);
1357+
await ingestDraftLaneSegments(payload, true);
13661358
})
13671359
: undefined,
13681360
onAssistantMessageStart: answerLane.stream

‎src/agents/pi-embedded-runner/run/params.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
import type { AgentMessage } from "@mariozechner/pi-agent-core";
22
import type { ImageContent } from "@mariozechner/pi-ai";
3-
import type { SourceReplyDeliveryMode } from "../../../auto-reply/get-reply-options.types.js";
3+
import type {
4+
PartialReplyPayload,
5+
SourceReplyDeliveryMode,
6+
} from "../../../auto-reply/get-reply-options.types.js";
47
import type { ReplyPayload } from "../../../auto-reply/reply-payload.js";
58
import type { ReplyOperation } from "../../../auto-reply/reply/reply-run-registry.js";
69
import type { ReasoningLevel, ThinkLevel, VerboseLevel } from "../../../auto-reply/thinking.js";
@@ -151,7 +154,7 @@ export type RunEmbeddedPiAgentParams = {
151154
replyOperation?: ReplyOperation;
152155
shouldEmitToolResult?: () => boolean;
153156
shouldEmitToolOutput?: () => boolean;
154-
onPartialReply?: (payload: { text?: string; mediaUrls?: string[] }) => void | Promise<void>;
157+
onPartialReply?: (payload: PartialReplyPayload) => void | Promise<void>;
155158
onAssistantMessageStart?: () => void | Promise<void>;
156159
onBlockReply?: (payload: BlockReplyPayload) => void | Promise<void>;
157160
onBlockReplyFlush?: () => void | Promise<void>;

‎src/agents/pi-embedded-subscribe.types.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { AgentSession } from "@mariozechner/pi-coding-agent";
2+
import type { PartialReplyPayload } from "../auto-reply/get-reply-options.types.js";
23
import type { ReplyPayload } from "../auto-reply/reply-payload.js";
34
import type { ReasoningLevel, ThinkLevel, VerboseLevel } from "../auto-reply/thinking.js";
45
import type { OpenClawConfig } from "../config/types.openclaw.js";
@@ -38,7 +39,7 @@ export type SubscribeEmbeddedPiSessionParams = {
3839
onBlockReplyFlush?: () => void | Promise<void>;
3940
blockReplyBreak?: "text_end" | "message_end";
4041
blockReplyChunking?: BlockReplyChunking;
41-
onPartialReply?: (payload: { text?: string; mediaUrls?: string[] }) => void | Promise<void>;
42+
onPartialReply?: (payload: PartialReplyPayload) => void | Promise<void>;
4243
onAssistantMessageStart?: () => void | Promise<void>;
4344
onAgentEvent?: (evt: {
4445
stream: string;

‎src/auto-reply/get-reply-options.types.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@ export type ReplyThreadingPolicy = {
3131

3232
export type SourceReplyDeliveryMode = "automatic" | "message_tool_only";
3333

34+
export type PartialReplyPayload = Pick<ReplyPayload, "text" | "mediaUrls"> & {
35+
delta?: string;
36+
replace?: true;
37+
};
38+
3439
export type GetReplyOptions = {
3540
/** Override run id for agent events (defaults to random UUID). */
3641
runId?: string;
@@ -72,7 +77,7 @@ export type GetReplyOptions = {
7277
* channel to surface progress via its own streaming/edit UX.
7378
*/
7479
suppressDefaultToolProgressMessages?: boolean;
75-
onPartialReply?: (payload: ReplyPayload) => Promise<void> | void;
80+
onPartialReply?: (payload: PartialReplyPayload) => Promise<void> | void;
7681
onReasoningStream?: (payload: ReplyPayload) => Promise<void> | void;
7782
/** Called when a thinking/reasoning block ends. */
7883
onReasoningEnd?: () => Promise<void> | void;

‎src/auto-reply/types.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
export type {
22
BlockReplyContext,
33
GetReplyOptions,
4+
PartialReplyPayload,
45
ReplyThreadingPolicy,
56
TypingPolicy,
67
} from "./get-reply-options.types.js";

0 commit comments

Comments
 (0)