483 lines
18 KiB
TypeScript
483 lines
18 KiB
TypeScript
import path from "node:path";
|
|
import type { Bot } from "grammy";
|
|
import { beforeEach, describe, expect, it, vi } from "vitest";
|
|
import { STATE_DIR } from "../config/paths.js";
|
|
|
|
const createTelegramDraftStream = vi.hoisted(() => vi.fn());
|
|
const dispatchReplyWithBufferedBlockDispatcher = vi.hoisted(() => vi.fn());
|
|
const deliverReplies = vi.hoisted(() => vi.fn());
|
|
const editMessageTelegram = vi.hoisted(() => vi.fn());
|
|
|
|
vi.mock("./draft-stream.js", () => ({
|
|
createTelegramDraftStream,
|
|
}));
|
|
|
|
vi.mock("../auto-reply/reply/provider-dispatcher.js", () => ({
|
|
dispatchReplyWithBufferedBlockDispatcher,
|
|
}));
|
|
|
|
vi.mock("./bot/delivery.js", () => ({
|
|
deliverReplies,
|
|
}));
|
|
|
|
vi.mock("./send.js", () => ({
|
|
editMessageTelegram,
|
|
}));
|
|
|
|
vi.mock("./sticker-cache.js", () => ({
|
|
cacheSticker: vi.fn(),
|
|
describeStickerImage: vi.fn(),
|
|
}));
|
|
|
|
import { dispatchTelegramMessage } from "./bot-message-dispatch.js";
|
|
|
|
describe("dispatchTelegramMessage draft streaming", () => {
|
|
type TelegramMessageContext = Parameters<typeof dispatchTelegramMessage>[0]["context"];
|
|
|
|
beforeEach(() => {
|
|
createTelegramDraftStream.mockReset();
|
|
dispatchReplyWithBufferedBlockDispatcher.mockReset();
|
|
deliverReplies.mockReset();
|
|
editMessageTelegram.mockReset();
|
|
});
|
|
|
|
function createDraftStream(messageId?: number) {
|
|
return {
|
|
update: vi.fn(),
|
|
flush: vi.fn().mockResolvedValue(undefined),
|
|
messageId: vi.fn().mockReturnValue(messageId),
|
|
clear: vi.fn().mockResolvedValue(undefined),
|
|
stop: vi.fn().mockResolvedValue(undefined),
|
|
forceNewMessage: vi.fn(),
|
|
};
|
|
}
|
|
|
|
function createContext(overrides?: Partial<TelegramMessageContext>): TelegramMessageContext {
|
|
const base = {
|
|
ctxPayload: {},
|
|
primaryCtx: { message: { chat: { id: 123, type: "private" } } },
|
|
msg: {
|
|
chat: { id: 123, type: "private" },
|
|
message_id: 456,
|
|
message_thread_id: 777,
|
|
},
|
|
chatId: 123,
|
|
isGroup: false,
|
|
resolvedThreadId: undefined,
|
|
replyThreadId: 777,
|
|
threadSpec: { id: 777, scope: "dm" },
|
|
historyKey: undefined,
|
|
historyLimit: 0,
|
|
groupHistories: new Map(),
|
|
route: { agentId: "default", accountId: "default" },
|
|
skillFilter: undefined,
|
|
sendTyping: vi.fn(),
|
|
sendRecordVoice: vi.fn(),
|
|
ackReactionPromise: null,
|
|
reactionApi: null,
|
|
removeAckAfterReply: false,
|
|
} as unknown as TelegramMessageContext;
|
|
|
|
return {
|
|
...base,
|
|
...overrides,
|
|
// Merge nested fields when overrides provide partial objects.
|
|
primaryCtx: {
|
|
...(base.primaryCtx as object),
|
|
...(overrides?.primaryCtx ? (overrides.primaryCtx as object) : null),
|
|
} as TelegramMessageContext["primaryCtx"],
|
|
msg: {
|
|
...(base.msg as object),
|
|
...(overrides?.msg ? (overrides.msg as object) : null),
|
|
} as TelegramMessageContext["msg"],
|
|
route: {
|
|
...(base.route as object),
|
|
...(overrides?.route ? (overrides.route as object) : null),
|
|
} as TelegramMessageContext["route"],
|
|
};
|
|
}
|
|
|
|
function createBot(): Bot {
|
|
return { api: { sendMessage: vi.fn(), editMessageText: vi.fn() } } as unknown as Bot;
|
|
}
|
|
|
|
function createRuntime(): Parameters<typeof dispatchTelegramMessage>[0]["runtime"] {
|
|
return {
|
|
log: vi.fn(),
|
|
error: vi.fn(),
|
|
exit: () => {
|
|
throw new Error("exit");
|
|
},
|
|
};
|
|
}
|
|
|
|
async function dispatchWithContext(params: {
|
|
context: TelegramMessageContext;
|
|
telegramCfg?: Parameters<typeof dispatchTelegramMessage>[0]["telegramCfg"];
|
|
streamMode?: Parameters<typeof dispatchTelegramMessage>[0]["streamMode"];
|
|
}) {
|
|
await dispatchTelegramMessage({
|
|
context: params.context,
|
|
bot: createBot(),
|
|
cfg: {},
|
|
runtime: createRuntime(),
|
|
replyToMode: "first",
|
|
streamMode: params.streamMode ?? "partial",
|
|
textLimit: 4096,
|
|
telegramCfg: params.telegramCfg ?? {},
|
|
opts: { token: "token" },
|
|
});
|
|
}
|
|
|
|
it("streams drafts in private threads and forwards thread id", async () => {
|
|
const draftStream = createDraftStream();
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
await replyOptions?.onPartialReply?.({ text: "Hello" });
|
|
await dispatcherOptions.deliver({ text: "Hello" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
const context = createContext({
|
|
route: {
|
|
agentId: "work",
|
|
} as unknown as TelegramMessageContext["route"],
|
|
});
|
|
await dispatchWithContext({ context });
|
|
|
|
expect(createTelegramDraftStream).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
chatId: 123,
|
|
thread: { id: 777, scope: "dm" },
|
|
}),
|
|
);
|
|
expect(draftStream.update).toHaveBeenCalledWith("Hello");
|
|
expect(deliverReplies).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
thread: { id: 777, scope: "dm" },
|
|
mediaLocalRoots: expect.arrayContaining([path.join(STATE_DIR, "workspace-work")]),
|
|
}),
|
|
);
|
|
expect(dispatchReplyWithBufferedBlockDispatcher).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
replyOptions: expect.objectContaining({
|
|
disableBlockStreaming: true,
|
|
}),
|
|
}),
|
|
);
|
|
expect(editMessageTelegram).not.toHaveBeenCalled();
|
|
expect(draftStream.clear).toHaveBeenCalledTimes(1);
|
|
});
|
|
|
|
it("keeps block streaming enabled when account config enables it", async () => {
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => {
|
|
await dispatcherOptions.deliver({ text: "Hello" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
});
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({
|
|
context: createContext(),
|
|
telegramCfg: { blockStreaming: true },
|
|
});
|
|
|
|
expect(createTelegramDraftStream).not.toHaveBeenCalled();
|
|
expect(dispatchReplyWithBufferedBlockDispatcher).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
replyOptions: expect.objectContaining({
|
|
disableBlockStreaming: false,
|
|
onPartialReply: undefined,
|
|
}),
|
|
}),
|
|
);
|
|
});
|
|
|
|
it("finalizes text-only replies by editing the preview message in place", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
await replyOptions?.onPartialReply?.({ text: "Hel" });
|
|
await dispatcherOptions.deliver({ text: "Hello final" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "999" });
|
|
|
|
await dispatchWithContext({ context: createContext() });
|
|
|
|
expect(editMessageTelegram).toHaveBeenCalledWith(123, 999, "Hello final", expect.any(Object));
|
|
expect(deliverReplies).not.toHaveBeenCalled();
|
|
expect(draftStream.clear).not.toHaveBeenCalled();
|
|
expect(draftStream.stop).toHaveBeenCalled();
|
|
});
|
|
|
|
it("edits the preview message created during stop() final flush", async () => {
|
|
let messageId: number | undefined;
|
|
const draftStream = {
|
|
update: vi.fn(),
|
|
flush: vi.fn().mockResolvedValue(undefined),
|
|
messageId: vi.fn().mockImplementation(() => messageId),
|
|
clear: vi.fn().mockResolvedValue(undefined),
|
|
stop: vi.fn().mockImplementation(async () => {
|
|
messageId = 777;
|
|
}),
|
|
forceNewMessage: vi.fn(),
|
|
};
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => {
|
|
await dispatcherOptions.deliver({ text: "Short final" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
});
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "777" });
|
|
|
|
await dispatchWithContext({ context: createContext() });
|
|
|
|
expect(editMessageTelegram).toHaveBeenCalledWith(123, 777, "Short final", expect.any(Object));
|
|
expect(deliverReplies).not.toHaveBeenCalled();
|
|
expect(draftStream.stop).toHaveBeenCalled();
|
|
});
|
|
|
|
it("does not overwrite finalized preview when additional final payloads are sent", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => {
|
|
await dispatcherOptions.deliver({ text: "Primary result" }, { kind: "final" });
|
|
await dispatcherOptions.deliver(
|
|
{ text: "⚠️ Recovered tool error details" },
|
|
{ kind: "final" },
|
|
);
|
|
return { queuedFinal: true };
|
|
});
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "999" });
|
|
|
|
await dispatchWithContext({ context: createContext() });
|
|
|
|
expect(editMessageTelegram).toHaveBeenCalledTimes(1);
|
|
expect(editMessageTelegram).toHaveBeenCalledWith(
|
|
123,
|
|
999,
|
|
"Primary result",
|
|
expect.any(Object),
|
|
);
|
|
expect(deliverReplies).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
replies: [expect.objectContaining({ text: "⚠️ Recovered tool error details" })],
|
|
}),
|
|
);
|
|
expect(draftStream.clear).not.toHaveBeenCalled();
|
|
expect(draftStream.stop).toHaveBeenCalled();
|
|
});
|
|
|
|
it("falls back to normal delivery when preview final is too long to edit", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
const longText = "x".repeat(5000);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => {
|
|
await dispatcherOptions.deliver({ text: longText }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
});
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
editMessageTelegram.mockResolvedValue({ ok: true, chatId: "123", messageId: "999" });
|
|
|
|
await dispatchWithContext({ context: createContext() });
|
|
|
|
expect(editMessageTelegram).not.toHaveBeenCalled();
|
|
expect(deliverReplies).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
replies: [expect.objectContaining({ text: longText })],
|
|
}),
|
|
);
|
|
expect(draftStream.clear).toHaveBeenCalledTimes(1);
|
|
expect(draftStream.stop).toHaveBeenCalled();
|
|
});
|
|
|
|
it("disables block streaming when streamMode is off", async () => {
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => {
|
|
await dispatcherOptions.deliver({ text: "Hello" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
});
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({
|
|
context: createContext(),
|
|
streamMode: "off",
|
|
});
|
|
|
|
expect(createTelegramDraftStream).not.toHaveBeenCalled();
|
|
expect(dispatchReplyWithBufferedBlockDispatcher).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
replyOptions: expect.objectContaining({
|
|
disableBlockStreaming: true,
|
|
}),
|
|
}),
|
|
);
|
|
});
|
|
|
|
it("forces new message when new assistant message starts after previous output", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
// First assistant message: partial text
|
|
await replyOptions?.onPartialReply?.({ text: "First response" });
|
|
// New assistant message starts (e.g., after tool call)
|
|
await replyOptions?.onAssistantMessageStart?.();
|
|
// Second assistant message: new text
|
|
await replyOptions?.onPartialReply?.({ text: "After tool call" });
|
|
await dispatcherOptions.deliver({ text: "After tool call" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({ context: createContext(), streamMode: "block" });
|
|
|
|
// Should force new message when assistant message starts after previous output
|
|
expect(draftStream.forceNewMessage).toHaveBeenCalled();
|
|
});
|
|
|
|
it("does not force new message in partial mode when assistant message restarts", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
await replyOptions?.onPartialReply?.({ text: "First response" });
|
|
await replyOptions?.onAssistantMessageStart?.();
|
|
await replyOptions?.onPartialReply?.({ text: "After tool call" });
|
|
await dispatcherOptions.deliver({ text: "After tool call" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({ context: createContext(), streamMode: "partial" });
|
|
|
|
expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("does not force new message on first assistant message start", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
// First assistant message starts (no previous output)
|
|
await replyOptions?.onAssistantMessageStart?.();
|
|
// Partial updates
|
|
await replyOptions?.onPartialReply?.({ text: "Hello" });
|
|
await replyOptions?.onPartialReply?.({ text: "Hello world" });
|
|
await dispatcherOptions.deliver({ text: "Hello world" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({ context: createContext(), streamMode: "block" });
|
|
|
|
// First message start shouldn't trigger forceNewMessage (no previous output)
|
|
expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("forces new message when reasoning ends after previous output", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
// First partial: text before thinking
|
|
await replyOptions?.onPartialReply?.({ text: "Let me check" });
|
|
// Reasoning stream (thinking block)
|
|
await replyOptions?.onReasoningStream?.({ text: "Analyzing..." });
|
|
// Reasoning ends
|
|
await replyOptions?.onReasoningEnd?.();
|
|
// Second partial: text after thinking
|
|
await replyOptions?.onPartialReply?.({ text: "Here's the answer" });
|
|
await dispatcherOptions.deliver({ text: "Here's the answer" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({ context: createContext(), streamMode: "block" });
|
|
|
|
// Should force new message when reasoning ends
|
|
expect(draftStream.forceNewMessage).toHaveBeenCalled();
|
|
});
|
|
|
|
it("does not force new message in partial mode when reasoning ends", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
await replyOptions?.onPartialReply?.({ text: "Let me check" });
|
|
await replyOptions?.onReasoningEnd?.();
|
|
await replyOptions?.onPartialReply?.({ text: "Here's the answer" });
|
|
await dispatcherOptions.deliver({ text: "Here's the answer" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({ context: createContext(), streamMode: "partial" });
|
|
|
|
expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("does not force new message on reasoning end without previous output", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
// Reasoning starts immediately (no previous text output)
|
|
await replyOptions?.onReasoningStream?.({ text: "Thinking..." });
|
|
// Reasoning ends
|
|
await replyOptions?.onReasoningEnd?.();
|
|
// First actual text output
|
|
await replyOptions?.onPartialReply?.({ text: "Here's my answer" });
|
|
await dispatcherOptions.deliver({ text: "Here's my answer" }, { kind: "final" });
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({ context: createContext(), streamMode: "block" });
|
|
|
|
// No previous text output, so no forceNewMessage needed
|
|
expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("does not edit preview message when final payload is an error", async () => {
|
|
const draftStream = createDraftStream(999);
|
|
createTelegramDraftStream.mockReturnValue(draftStream);
|
|
dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
|
|
async ({ dispatcherOptions, replyOptions }) => {
|
|
// Partial text output
|
|
await replyOptions?.onPartialReply?.({ text: "Let me check that file" });
|
|
// Error payload should not edit the preview message
|
|
await dispatcherOptions.deliver(
|
|
{ text: "⚠️ 🛠️ Exec: cat /nonexistent failed: No such file", isError: true },
|
|
{ kind: "final" },
|
|
);
|
|
return { queuedFinal: true };
|
|
},
|
|
);
|
|
deliverReplies.mockResolvedValue({ delivered: true });
|
|
|
|
await dispatchWithContext({ context: createContext(), streamMode: "block" });
|
|
|
|
// Should NOT edit preview message (which would overwrite the partial text)
|
|
expect(editMessageTelegram).not.toHaveBeenCalled();
|
|
// Should deliver via normal path as a new message
|
|
expect(deliverReplies).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
replies: [expect.objectContaining({ text: expect.stringContaining("⚠️") })],
|
|
}),
|
|
);
|
|
});
|
|
});
|