From 8281904cb23121f38b9e646072358c6d340027fd Mon Sep 17 00:00:00 2001 From: Gourav Date: Tue, 21 Apr 2026 18:33:53 +0530 Subject: [PATCH] feat(chat): enhance message streaming logic for tool responses - Introduced a mechanism to handle post-tool text streaming when a message:delta is received without a preceding message:start, allowing for live updates. - Added a guard to prevent duplicate insertion of assistant messages that follow tool results, ensuring proper message flow. - Improved debug logging for better traceability of streaming state changes. This update refines the chat experience by ensuring that tool responses are integrated seamlessly into the conversation flow. --- .../src/chat/classes/AbstractChat.ts | 62 +++++++++++++++---- 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/packages/copilot-sdk/src/chat/classes/AbstractChat.ts b/packages/copilot-sdk/src/chat/classes/AbstractChat.ts index da38799..351cecb 100644 --- a/packages/copilot-sdk/src/chat/classes/AbstractChat.ts +++ b/packages/copilot-sdk/src/chat/classes/AbstractChat.ts @@ -1087,6 +1087,10 @@ export class AbstractChat { let chunkCount = 0; let toolCallsEmitted = false; // Guard to prevent emitting toolCalls twice + // Set to true when post-tool text was already streamed live via auto-init + // from a message:delta (no message:start). Prevents done handler from + // inserting a duplicate when seenToolResult is true. + let postToolTextStreamed = false; // Holds client tool calls received via a tool_calls chunk AFTER a // mid-stream message:end nulled streamState. // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -1188,10 +1192,30 @@ export class AbstractChat { // Update stream state (pure function) // Skip most chunks if streamState is null. - // EXCEPTION: after a mid-stream message:end the server can still send + // EXCEPTION 1: after a mid-stream message:end the server can still send // tool_calls + done for client-side tool dispatch. Handle those directly. + // EXCEPTION 2: some server-side tools (e.g. websearch) stream the post-tool + // response via message:delta WITHOUT a preceding message:start. Auto-init a + // new message so the text streams live instead of appearing all-at-once from done. if (!this.streamState) { - if (chunk.type === "tool_calls") { + if (chunk.type === "message:delta") { + this.debug( + "message:delta with no streamState — auto-init new message", + ); + const currentLeaf = this.state.messages; + const currentLeafId = + currentLeaf.length > 0 + ? currentLeaf[currentLeaf.length - 1].id + : undefined; + const newMessage = createEmptyAssistantMessage(undefined, { + parentId: currentLeafId, + }) as T; + this.state.pushMessage(newMessage); + this.streamState = createStreamState(newMessage.id); + this.callbacks.onMessageStart?.(newMessage.id); + postToolTextStreamed = true; + // No continue — fall out of this block into normal processStreamChunk below + } else if (chunk.type === "tool_calls") { // Store for emission when done arrives. Do NOT update message state // here — done.messages carries the assistant message with tool_calls // in proper OpenAI format, which we use in the done handler below. @@ -1202,9 +1226,7 @@ export class AbstractChat { ids: pendingClientToolCalls?.map((tc: { id?: string }) => tc.id), }); continue; - } - - if (chunk.type === "done") { + } else if (chunk.type === "done") { this.debug("done (post-message:end)", { hasPendingToolCalls: !!pendingClientToolCalls?.length, pendingCount: pendingClientToolCalls?.length ?? 0, @@ -1340,10 +1362,12 @@ export class AbstractChat { }); } continue; + } else { + this.debug("warning", "streamState is null, skipping chunk"); + continue; } - - this.debug("warning", "streamState is null, skipping chunk"); - continue; + // Only message:delta reaches here — streamState was just auto-init'd. + // Fall through to processStreamChunk below. } this.streamState = processStreamChunk(chunk, this.streamState); @@ -1474,13 +1498,25 @@ export class AbstractChat { } } + // Track whether a tool-result message has been seen as we iterate done.messages. + // Assistant text that appears AFTER a tool result was never streamed live + // (server-side builtin tools like websearch/webanswer return the full history + // only in the done payload). Assistant text that appears BEFORE any tool result + // is the intro turn already represented by streamed message:start/delta/end events. + let seenToolResult = false; for (const msg of chunk.messages) { - // Skip plain assistant text messages because they are already represented - // by streamed message:start/message:delta/message:end events. Preserve - // assistant messages that carry tool_calls so tool results keep a valid - // preceding assistant tool_call message in local state. + if (msg.role === "tool") seenToolResult = true; + + // Skip plain assistant text messages that precede any tool result — those + // are already represented by streamed message:start/message:delta/message:end + // events. Assistant messages that appear AFTER a tool result are post-tool + // responses that were never streamed live and must be inserted — UNLESS + // postToolTextStreamed is true, meaning message:delta already delivered them + // live via the auto-init path (skipping avoids a duplicate message). if (msg.role === "assistant" && !msg.tool_calls?.length) { - continue; + if (!seenToolResult) continue; + if (postToolTextStreamed) continue; + // Post-tool text: fall through to insert. } // The current streamed turn already becomes an assistant message from