diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts b/src/agents/embedded-agent-subscribe.handlers.messages.test.ts index dab3a1320d7e..470805bcb628 100644 --- a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts +++ b/src/agents/embedded-agent-subscribe.handlers.messages.test.ts @@ -478,6 +478,141 @@ describe("handleMessageUpdate text signatures", () => { ]); }); + it("uses incremental deltas for same-item phased streams", () => { + const onAgentEvent = vi.fn(); + const context = createMessageUpdateContext({ onAgentEvent }); + const signature = JSON.stringify({ v: 1, id: "item-final", phase: "final_answer" }); + const partial = { + role: "assistant", + phase: "final_answer", + content: [ + { + type: "text", + textSignature: signature, + get text() { + throw new Error("full partial text should not be read"); + }, + }, + ], + }; + + const createPhasedDelta = (delta: string) => + ({ + type: "message_update", + message: { role: "assistant", content: [] }, + assistantMessageEvent: { + type: "text_delta", + delta, + partial, + }, + }) as never; + + handleMessageUpdate(context, createPhasedDelta("Hello")); + handleMessageUpdate(context, createPhasedDelta(" world")); + + expect(onAgentEvent.mock.calls.map(([event]) => event)).toMatchObject([ + { + stream: "assistant", + data: { text: "Hello", delta: "Hello", phase: "final_answer" }, + }, + { + stream: "assistant", + data: { text: "Hello world", delta: " world", phase: "final_answer" }, + }, + ]); + }); + + it("keeps same-item phased stream deltas on the user-visible sanitizer path", () => { + const onAgentEvent = vi.fn(); + const context = createMessageUpdateContext({ onAgentEvent }); + const signature = JSON.stringify({ v: 1, id: "item-final", phase: "final_answer" }); + const partial = { + role: "assistant", + phase: "final_answer", + content: [ + { + type: "text", + textSignature: signature, + get text() { + throw new Error("full partial text should not be read"); + }, + }, + ], + }; + + const createPhasedDelta = (delta: string) => + ({ + type: "message_update", + message: { role: "assistant", content: [] }, + assistantMessageEvent: { + type: "text_delta", + delta, + partial, + }, + }) as never; + + handleMessageUpdate(context, createPhasedDelta("Visible\n{")); + handleMessageUpdate( + context, + createPhasedDelta('"name":"read","arguments":{"file_path":"secret.md"}}'), + ); + handleMessageUpdate(context, createPhasedDelta("\nDone.")); + + expect(onAgentEvent.mock.calls.map(([event]) => event)).toMatchObject([ + { + stream: "assistant", + data: { text: "Visible", delta: "Visible", phase: "final_answer" }, + }, + { + stream: "assistant", + data: { text: "Visible\n\nDone.", delta: "\n\nDone.", phase: "final_answer" }, + }, + ]); + }); + + it("keeps sanitizer context when a same-item phased stream starts hidden", () => { + const onAgentEvent = vi.fn(); + const context = createMessageUpdateContext({ onAgentEvent }); + const signature = JSON.stringify({ v: 1, id: "item-final", phase: "final_answer" }); + const partial = { + role: "assistant", + phase: "final_answer", + content: [ + { + type: "text", + textSignature: signature, + get text() { + throw new Error("full partial text should not be read"); + }, + }, + ], + }; + + const createPhasedDelta = (delta: string) => + ({ + type: "message_update", + message: { role: "assistant", content: [] }, + assistantMessageEvent: { + type: "text_delta", + delta, + partial, + }, + }) as never; + + handleMessageUpdate(context, createPhasedDelta("{")); + handleMessageUpdate( + context, + createPhasedDelta('"name":"read","arguments":{"file_path":"secret.md"}}\nDone.'), + ); + + expect(onAgentEvent.mock.calls.map(([event]) => event)).toMatchObject([ + { + stream: "assistant", + data: { text: "Done.", delta: "Done.", phase: "final_answer" }, + }, + ]); + }); + it("treats phased textSignature item changes as assistant-message boundaries", () => { const flushBlockReplyBuffer = vi.fn(); const resetAssistantMessageState = vi.fn(); diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.ts b/src/agents/embedded-agent-subscribe.handlers.messages.ts index c96da4e97645..d4deb74a9f90 100644 --- a/src/agents/embedded-agent-subscribe.handlers.messages.ts +++ b/src/agents/embedded-agent-subscribe.handlers.messages.ts @@ -38,6 +38,7 @@ import { extractThinkingFromTaggedStream, extractThinkingFromTaggedText, promoteThinkingTagsToBlocks, + sanitizeAssistantVisibleStreamText, } from "./embedded-agent-utils.js"; import type { AgentEvent, AgentMessage } from "./runtime/index.js"; @@ -218,6 +219,22 @@ function resolveStreamVisibleText(params: { return { rawText, visibleText: rawText.trim() }; } +function resolveTextAppendDelta(previousText: string, nextText: string): string { + if (!nextText) { + return ""; + } + if (!previousText) { + return nextText; + } + if (nextText.startsWith(previousText)) { + return nextText.slice(previousText.length); + } + if (previousText.startsWith(nextText)) { + return ""; + } + return nextText; +} + function copyPartialBlockState( target: EmbeddedAgentSubscribeState["partialBlockState"], source: EmbeddedAgentSubscribeState["partialBlockState"], @@ -647,9 +664,11 @@ export function handleMessageUpdate( !deliveryPhase && Boolean(streamItemId) && isOpenAiResponsesAssistantMessage(partialAssistant); + let streamItemChanged = false; if ((deliveryPhase || isPhasePendingOpenAiResponsesTextItem) && streamItemId) { const previousStreamItemId = ctx.state.lastAssistantStreamItemId; if (previousStreamItemId && previousStreamItemId !== streamItemId) { + streamItemChanged = true; void ctx.flushBlockReplyBuffer({ assistantMessageIndex: ctx.state.assistantMessageIndex }); ctx.resetAssistantMessageState(ctx.state.assistantTexts.length); void ctx.params.onAssistantMessageStart?.(); @@ -677,11 +696,29 @@ export function handleMessageUpdate( } const wasThinking = ctx.state.partialBlockState.thinking; let visibleDelta = ""; - let next = shouldUsePhaseAwareBlockReply + const shouldReadPhaseAwarePartialText = + shouldUsePhaseAwareBlockReply && (streamItemChanged || evtType === "text_end" || !chunk); + let next = shouldReadPhaseAwarePartialText ? coerceChatContentText(extractAssistantVisibleText(partialAssistant)).trim() : ""; let nextRawStreamText = next; - if (!next && deliveryPhase !== "final_answer") { + let shouldPersistRawStreamText = false; + if (shouldUsePhaseAwareBlockReply && !next && deliveryPhase === "final_answer" && chunk) { + visibleDelta = ctx.stripBlockTags(chunk, ctx.state.partialBlockState, { + final: evtType === "text_end", + }); + const streamVisibleText = resolveStreamVisibleText({ + previousRawText: ctx.state.lastStreamedAssistant ?? "", + visibleDelta, + }); + const previousVisibleText = sanitizeAssistantVisibleStreamText( + ctx.state.lastStreamedAssistant ?? "", + ).trim(); + next = sanitizeAssistantVisibleStreamText(streamVisibleText.rawText).trim(); + visibleDelta = resolveTextAppendDelta(previousVisibleText, next); + nextRawStreamText = streamVisibleText.rawText; + shouldPersistRawStreamText = true; + } else if (!next && deliveryPhase !== "final_answer") { const pendingTagFragment = ctx.state.partialBlockState.pendingTagFragment; const shouldRecomputeFullStream = Boolean(pendingTagFragment) || REASONING_TAG_RE.test(chunk); if (shouldRecomputeFullStream) { @@ -800,6 +837,8 @@ export function handleMessageUpdate( ctx.emitAssistantStreamData(data, { emitPartialReply: true }); ctx.state.emittedAssistantUpdate = true; } + } else if (shouldPersistRawStreamText) { + ctx.state.lastStreamedAssistant = nextRawStreamText; } if ( diff --git a/src/agents/embedded-agent-utils.ts b/src/agents/embedded-agent-utils.ts index a3672c5a1aeb..df14ce57bbe7 100644 --- a/src/agents/embedded-agent-utils.ts +++ b/src/agents/embedded-agent-utils.ts @@ -40,6 +40,10 @@ function sanitizeAssistantText(text: string): string { return sanitizeAssistantVisibleText(text); } +export function sanitizeAssistantVisibleStreamText(text: string): string { + return sanitizeUserFacingText(sanitizeAssistantText(text), { errorContext: false }); +} + function finalizeAssistantExtraction(msg: AssistantMessage, extracted: string): string { const errorContext = msg.stopReason === "error"; return sanitizeUserFacingText(extracted, { errorContext });