diff --git a/apps/web/src/domains/runtime/session-stream/session-stream-render-scheduler.ts b/apps/web/src/domains/runtime/session-stream/session-stream-render-scheduler.ts index 243ef132..d166e81f 100644 --- a/apps/web/src/domains/runtime/session-stream/session-stream-render-scheduler.ts +++ b/apps/web/src/domains/runtime/session-stream/session-stream-render-scheduler.ts @@ -70,8 +70,15 @@ const segmentGraphemes = createGraphemeSegmenter(); // Only visible streaming text is paced. Tool-call args and lifecycle events // pass through so pacing never delays structural updates, and reasoning-chunk / // thinking deltas are live-state reducer no-ops not worth throttling. -function getPaceableTextDelta(event: AgUiSessionEvent): string | null { +function getPaceableTextDelta( + event: AgUiSessionEvent, + messageRole: "assistant" | "user" | undefined, +): string | null { if (event.type === "REASONING_MESSAGE_CONTENT" || event.type === "TEXT_MESSAGE_CONTENT") { + if (event.type === "TEXT_MESSAGE_CONTENT" && messageRole === "user") { + return null; + } + return event.delta.length > 0 ? event.delta : null; } @@ -150,6 +157,7 @@ export class SessionStreamRenderScheduler { #frameHandle: number | null = null; readonly #host: SessionStreamRenderSchedulerHost; #lastDrainAt: number | null = null; + readonly #messageRoles = new Map(); #pacingCarry = 0; #queue: QueuedSessionEvent[] = []; #queueOffset = 0; @@ -168,6 +176,7 @@ export class SessionStreamRenderScheduler { this.#queue = []; this.#queueOffset = 0; this.#lastDrainAt = null; + this.#messageRoles.clear(); this.#pacingCarry = 0; } @@ -181,7 +190,23 @@ export class SessionStreamRenderScheduler { } for (const event of events) { + const messageKey = + "messageId" in event && typeof event.messageId === "string" + ? `${sessionId}\0${event.messageId}` + : null; + + if ( + event.type === "TEXT_MESSAGE_START" && + (event.role === "assistant" || event.role === "user") + ) { + this.#messageRoles.set(`${sessionId}\0${event.messageId}`, event.role); + } + this.#queue.push(this.#toQueueItem(sessionId, event)); + + if (event.type === "TEXT_MESSAGE_END" && messageKey !== null) { + this.#messageRoles.delete(messageKey); + } } this.#schedule(); @@ -411,7 +436,11 @@ export class SessionStreamRenderScheduler { } #toQueueItem(sessionId: string, event: AgUiSessionEvent): QueuedSessionEvent { - const delta = getPaceableTextDelta(event); + const messageRole = + "messageId" in event && typeof event.messageId === "string" + ? this.#messageRoles.get(`${sessionId}\0${event.messageId}`) + : undefined; + const delta = getPaceableTextDelta(event, messageRole); return { event, diff --git a/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts b/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts index 4d01ef65..6ad0d016 100644 --- a/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts +++ b/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts @@ -109,9 +109,16 @@ function sessionSegmentsToParts(segments: readonly SessionViewSegment[]): Assist // keeps assistant-ui's per-message memo + keys aligned with the live stream. // Permission/needs_approval is intentionally NOT folded here (this conversion is // identity-cached and permissionRequests live in a separate array). -export function convertSessionMessage(message: SessionViewMessage): ThreadMessageLike { +export function convertSessionMessage( + message: SessionViewMessage, + messageIndex?: number, +): ThreadMessageLike { if (message.role === "user") { - return { role: "user", id: message.id, content: message.content }; + // The optimistic message and its server echo have different source ids but + // occupy the same append-only transcript slot. Keep the UI id stable so + // assistant-ui updates the bubble instead of deleting and remounting it. + const id = messageIndex === undefined ? message.id : `user:${messageIndex}`; + return { role: "user", id, content: message.content }; } const parts = sessionSegmentsToParts(message.segments); diff --git a/apps/web/tests/session-message-convert.test.ts b/apps/web/tests/session-message-convert.test.ts index d844c2d8..6e5894ac 100644 --- a/apps/web/tests/session-message-convert.test.ts +++ b/apps/web/tests/session-message-convert.test.ts @@ -29,6 +29,34 @@ describe("convertSessionMessage", () => { expect(message.role).toBe("user"); }); + test("keeps a user bubble identity when the optimistic message is replaced by its echo", () => { + const optimistic = convertSessionMessage( + { + content: "hello", + createdAt: "", + id: "pending:request-1", + plan: [], + role: "user", + segments: [], + }, + 2, + ); + const echoed = convertSessionMessage( + { + content: "hello", + createdAt: "", + id: "01J0000000000000000000000U", + plan: [], + role: "user", + segments: [], + }, + 2, + ); + + expect(optimistic.id).toBe("user:2"); + expect(echoed.id).toBe(optimistic.id); + }); + test("merges consecutive text segments into one markdown part", () => { const like = convertSessionMessage( assistantMessage([ diff --git a/apps/web/tests/session-stream-render-scheduler.test.ts b/apps/web/tests/session-stream-render-scheduler.test.ts index c2d9f66f..d8c73035 100644 --- a/apps/web/tests/session-stream-render-scheduler.test.ts +++ b/apps/web/tests/session-stream-render-scheduler.test.ts @@ -354,6 +354,25 @@ describe("session stream render scheduler", () => { expect(drainFrames(manual)).toBe(0); }); + test("delivers user start/content events atomically instead of pacing the server echo", () => { + const manual = createManualHost(); + const batches: AgUiSessionEvent[][] = []; + const scheduler = new SessionStreamRenderScheduler((_sessionId, events) => { + batches.push(events); + return true; + }, manual.host); + const events = [ + { messageId: "message-user", role: "user", type: "TEXT_MESSAGE_START" }, + textEvent("用".repeat(1000), "message-user"), + ] as const satisfies readonly AgUiSessionEvent[]; + + scheduler.enqueueMany("session-1", [...events]); + manual.fireFrame(); + + expect(batches).toEqual([[...events]]); + expect(drainFrames(manual)).toBe(0); + }); + test("flushNow delivers the paced remainder exactly once", () => { const manual = createManualHost(); const batches: AgUiSessionEvent[][] = [];