diff --git a/src/channels/draft-stream.ts b/src/channels/draft-stream.ts index 75c3f0b..96f860a 100644 --- a/src/channels/draft-stream.ts +++ b/src/channels/draft-stream.ts @@ -7,15 +7,14 @@ import { DraftThrottle } from "./draft-throttle.js"; import { formatThinkingBlock, splitMessageText } from "./progress.js"; /** - * Channel-agnostic streaming reply: one editable progress message whose edits - * are throttled to one per interval, followed by the final answer as separate - * messages. Connectors supply the transport (post/update) and the + * Channel-agnostic streaming reply: one editable message that starts as + * progress and is replaced by the final answer. Edits are throttled to one + * per interval. Connectors supply the transport (post/update) and the * channel-specific text rendering; all draft/lifecycle state lives here. */ export abstract class DraftReplyStream implements OutboundStream { static readonly #draftIntervalMs = 1_500; #progress: ProgressSnapshot = { actions: [], plan: [] }; - #finalText = ""; #messageId: string | undefined; #starting: Promise | undefined; #lastPublishedText = ""; @@ -116,10 +115,11 @@ export abstract class DraftReplyStream implements OutboundStream { this.scheduleDraft(); } - public appendFinal(delta: string): void { - if (this.#lifecycle.closed) return; - this.#finalText += delta; - this.scheduleDraft(); + public appendFinal(_delta: string): void { + // Keep the editable message as an unambiguous progress placeholder until + // complete() can replace it with the authoritative final text. If a + // streamed preview were published first and the final edit failed, the + // delivery fallback would leave that stale preview beside the full answer. } public async complete( @@ -144,22 +144,26 @@ export abstract class DraftReplyStream implements OutboundStream { } const chunks = text.length === 0 ? [] : splitMessageText(this.renderFinal(text), this.#textLimit); - // Freeze the progress message without the streaming cursor. The answer - // itself arrives as separate messages below: a silent edit of the - // thinking message never notifies anyone, and a failed edit must not - // take the answer down with it. + const [first, ...remaining] = chunks; + let chunksToPost = chunks; if (this.#messageId !== undefined) { - await this.update( - this.#messageId, - this.renderProgress(formatThinkingBlock(this.#progress)).slice(0, this.#textLimit), - ).catch((error: unknown) => { - this.logger.debug(`${this.#channelLabel} progress freeze failed`, { - error: errorMessage(error), - }); - }); + const replacement = + first ?? this.renderProgress(formatThinkingBlock(this.#progress)).slice(0, this.#textLimit); + if (replacement !== this.#lastPublishedText) { + try { + await this.update(this.#messageId, replacement); + chunksToPost = remaining; + } catch (error) { + this.logger.debug(`${this.#channelLabel} final replacement failed`, { + error: errorMessage(error), + }); + } + } else { + chunksToPost = remaining; + } } let undelivered = 0; - for (const chunk of chunks) { + for (const chunk of chunksToPost) { try { await this.post(chunk); } catch (error) { @@ -197,16 +201,6 @@ export abstract class DraftReplyStream implements OutboundStream { } private preview(): string { - // Rendering can expand the progress block past the message limit, so - // clamp it before budgeting the final-text tail. - const progress = this.renderProgress(formatThinkingBlock(this.#progress)).slice( - 0, - this.#textLimit - 3, - ); - if (this.#finalText.length === 0) return `${progress}\n\n▌`; - // available === 0 must not fall through to slice(-0), which is slice(0). - const available = Math.max(0, this.#textLimit - progress.length - 3); - const finalText = available === 0 ? "" : this.renderFinal(this.#finalText).slice(-available); - return `${progress}\n\n${finalText}▌`; + return this.renderProgress(formatThinkingBlock(this.#progress)).slice(0, this.#textLimit); } } diff --git a/test/draft-stream.test.ts b/test/draft-stream.test.ts new file mode 100644 index 0000000..dc0b6d2 --- /dev/null +++ b/test/draft-stream.test.ts @@ -0,0 +1,90 @@ +import { describe, expect, it } from "bun:test"; +import { DraftReplyStream } from "../src/channels/draft-stream.js"; +import { Logger } from "../src/shared/logger.js"; + +class RecordingDraftStream extends DraftReplyStream { + public readonly posts: string[] = []; + public readonly updates: string[] = []; + public failUpdates = false; + + public constructor(limit = 100) { + super(new Logger("error"), "Test", limit); + } + + protected async postInitial(content: string): Promise { + this.posts.push(content); + return "message-1"; + } + + protected async post(content: string): Promise { + this.posts.push(content); + } + + protected async update(_messageId: string, content: string): Promise { + if (this.failUpdates) throw new Error("update failed"); + this.updates.push(content); + } + + protected renderProgress(block: string): string { + return block; + } + + protected renderFinal(text: string): string { + return text; + } +} + +describe("DraftReplyStream", () => { + it("posts one thinking block and replaces it with the final answer", async () => { + const stream = new RecordingDraftStream(); + + await stream.start({ summary: "Thinking…", actions: [], plan: [] }); + expect(stream.posts).toEqual(["▌ Thinking…"]); + + await stream.complete("Finished"); + expect(stream.updates).toEqual(["Finished"]); + expect(stream.posts).toEqual(["▌ Thinking…"]); + }); + + it("keeps the thinking placeholder until the complete answer is available", async () => { + const stream = new RecordingDraftStream(); + + await stream.start(); + stream.appendFinal("Working answer"); + + expect(stream.posts).toEqual(["▌ Thinking…"]); + expect(stream.updates).toEqual([]); + }); + + it("posts only overflow chunks after replacing the thinking message", async () => { + const stream = new RecordingDraftStream(10); + + await stream.start(); + await stream.complete("1234567890abcdefghij"); + + expect(stream.updates).toEqual(["1234567890"]); + expect(stream.posts.slice(1)).toEqual(["abcdefghij"]); + }); + + it("posts the complete answer when replacing the thinking message fails", async () => { + const stream = new RecordingDraftStream(); + + await stream.start(); + stream.appendFinal("Partial preview"); + stream.failUpdates = true; + await stream.complete("Finished"); + + expect(stream.updates).toEqual([]); + expect(stream.posts).toEqual(["▌ Thinking…", "Finished"]); + }); + + it("publishes the latest progress when a turn ends without an answer", async () => { + const stream = new RecordingDraftStream(); + + await stream.start({ summary: "Thinking…", actions: [], plan: [] }); + stream.setProgress({ summary: "Waiting for input", actions: [], plan: [] }); + await stream.complete(""); + + expect(stream.updates).toEqual(["▌ Waiting for input"]); + }); +});