Skip to content

Commit aac174b

Browse files
vincentkocvincentkoc
committed
fix(provider): scope text replacement clears by index
1 parent 8e92f92 commit aac174b

2 files changed

Lines changed: 91 additions & 12 deletions

File tree

‎src/plugin-sdk/provider-stream-shared.test.ts‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,71 @@ describe("createPlainTextToolCallCompatWrapper", () => {
158158
]);
159159
expect(events.at(-1)).toMatchObject({ type: "done", reason: "toolUse" });
160160
});
161+
162+
it("does not clear flushed text from a different content index", async () => {
163+
const toolCallText = '[tool:read] {"path":"README.md"}';
164+
const finalMessage: AssistantMessage = {
165+
role: "assistant",
166+
api: "openai-compatible",
167+
provider: "test-provider",
168+
model: "test-model",
169+
content: [{ type: "text", text: toolCallText }],
170+
stopReason: "stop",
171+
timestamp: 123,
172+
usage: {
173+
input: 0,
174+
output: 0,
175+
cacheRead: 0,
176+
cacheWrite: 0,
177+
totalTokens: 0,
178+
cost: {
179+
input: 0,
180+
output: 0,
181+
cacheRead: 0,
182+
cacheWrite: 0,
183+
total: 0,
184+
},
185+
},
186+
};
187+
const emptyPartial: AssistantMessage = { ...finalMessage, content: [] };
188+
const baseStreamFn: StreamFn = () =>
189+
({
190+
async *[Symbol.asyncIterator]() {
191+
yield {
192+
type: "text_delta",
193+
contentIndex: 0,
194+
delta: "I will check.",
195+
partial: emptyPartial,
196+
};
197+
yield {
198+
type: "text_delta",
199+
contentIndex: 1,
200+
delta: toolCallText,
201+
replace: true,
202+
partial: emptyPartial,
203+
};
204+
yield { type: "done", reason: "stop", message: finalMessage };
205+
},
206+
result: async () => finalMessage,
207+
}) as ReturnType<StreamFn>;
208+
209+
const wrapped = createPlainTextToolCallCompatWrapper(baseStreamFn);
210+
const stream = await Promise.resolve(
211+
wrapped({} as never, { tools: [{ name: "read" }] } as never, {}),
212+
);
213+
const events: Array<Record<string, unknown>> = [];
214+
for await (const event of stream as AsyncIterable<Record<string, unknown>>) {
215+
events.push(event);
216+
}
217+
218+
expect(events.filter((event) => event.type === "text_delta")).toMatchObject([
219+
{ contentIndex: 0, delta: "I will check." },
220+
]);
221+
expect(events.some((event) => event.type === "text_delta" && event.replace === true)).toBe(
222+
false,
223+
);
224+
expect(events.at(-1)).toMatchObject({ type: "done", reason: "toolUse" });
225+
});
161226
});
162227

163228
describe("createDeepSeekV4OpenAICompatibleThinkingWrapper", () => {

‎src/plugin-sdk/provider-stream-shared.ts‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,12 @@ function toRecord(value: unknown): Record<string, unknown> | undefined {
3434
return value && typeof value === "object" ? (value as Record<string, unknown>) : undefined;
3535
}
3636

37+
function resolveStreamContentIndex(record: Record<string, unknown>): number {
38+
return typeof record.contentIndex === "number" && Number.isInteger(record.contentIndex)
39+
? record.contentIndex
40+
: 0;
41+
}
42+
3743
function resolveContextToolNames(context: Parameters<StreamFn>[1]): Set<string> {
3844
const tools = (context as { tools?: unknown }).tools;
3945
if (!Array.isArray(tools)) {
@@ -323,8 +329,8 @@ function wrapPlainTextToolCallStream(
323329
void (async () => {
324330
const bufferedTextEvents: unknown[] = [];
325331
let bufferedText = "";
326-
let flushedTextEvents = false;
327-
let replacementSupersedesFlushedText = false;
332+
const flushedTextContentIndexes = new Set<number>();
333+
let replacementClearContentIndex: number | undefined;
328334
let ended = false;
329335
const endStream = () => {
330336
if (!ended) {
@@ -334,18 +340,25 @@ function wrapPlainTextToolCallStream(
334340
};
335341
const flushBufferedTextEvents = () => {
336342
const events = bufferedTextEvents.splice(0);
337-
if (events.length > 0) {
338-
flushedTextEvents = true;
339-
}
340343
for (const event of events) {
344+
const record = toRecord(event);
345+
if (
346+
record &&
347+
(record.type === "text_start" ||
348+
record.type === "text_delta" ||
349+
record.type === "text_end")
350+
) {
351+
flushedTextContentIndexes.add(resolveStreamContentIndex(record));
352+
}
341353
stream.push(event);
342354
}
343355
bufferedText = "";
356+
replacementClearContentIndex = undefined;
344357
};
345-
const emitTextReplacementClear = (message: Record<string, unknown>) => {
358+
const emitTextReplacementClear = (message: Record<string, unknown>, contentIndex: number) => {
346359
stream.push({
347360
type: "text_delta",
348-
contentIndex: 0,
361+
contentIndex,
349362
delta: "",
350363
replace: true,
351364
partial: { ...message, content: [] },
@@ -359,11 +372,12 @@ function wrapPlainTextToolCallStream(
359372

360373
if (type === "text_start" || type === "text_delta" || type === "text_end") {
361374
const replacesBufferedText = type === "text_delta" && record?.replace === true;
375+
const contentIndex = record ? resolveStreamContentIndex(record) : 0;
362376
if (replacesBufferedText) {
363377
bufferedTextEvents.splice(0);
364-
if (flushedTextEvents) {
365-
replacementSupersedesFlushedText = true;
366-
}
378+
replacementClearContentIndex = flushedTextContentIndexes.has(contentIndex)
379+
? contentIndex
380+
: undefined;
367381
}
368382
bufferedTextEvents.push(event);
369383
if (typeof record?.delta === "string") {
@@ -387,8 +401,8 @@ function wrapPlainTextToolCallStream(
387401
if (promotedMessage) {
388402
bufferedTextEvents.splice(0);
389403
bufferedText = "";
390-
if (replacementSupersedesFlushedText) {
391-
emitTextReplacementClear(promotedMessage);
404+
if (replacementClearContentIndex !== undefined) {
405+
emitTextReplacementClear(promotedMessage, replacementClearContentIndex);
392406
}
393407
emitPromotedToolCallEvents(stream, promotedMessage);
394408
stream.push({ ...record, reason: "toolUse", message: promotedMessage });

0 commit comments

Comments
 (0)