diff --git a/packages/app/src/agent-stream/layout.test.ts b/packages/app/src/agent-stream/layout.test.ts index b93ef7acd..b996543be 100644 --- a/packages/app/src/agent-stream/layout.test.ts +++ b/packages/app/src/agent-stream/layout.test.ts @@ -112,6 +112,26 @@ function footerOwners(layout: StreamLayout): string[] { return owners; } +function footerAssistantIds(layout: StreamLayout): string[] { + return [ + ...layout.history.flatMap((item) => + item.completedFooter ? [item.completedFooter.itemId] : [], + ), + ...layout.liveHead.flatMap((item) => + item.completedFooter ? [item.completedFooter.itemId] : [], + ), + ...(layout.auxiliaryTurnFooter ? [layout.auxiliaryTurnFooter.itemId] : []), + ]; +} + +function inlineFooterPlacementByItemId(layout: StreamLayout): Record { + return Object.fromEntries( + [...layout.history, ...layout.liveHead].flatMap((item) => + item.completedFooter ? [[item.item.id, item.completedFooter.itemId]] : [], + ), + ); +} + function findLayoutItem(layout: StreamLayout, id: string): StreamLayoutItem { const item = [...layout.history, ...layout.liveHead].find( (candidate) => candidate.item.id === id, @@ -291,7 +311,7 @@ describe("layoutStream", () => { }); it.each(["web", "android"] as const)( - "keeps inline footer on an assistant turn with trailing tool rows before the next user on %s", + "places inline footer after trailing visible tool rows before the next user on %s", (platform) => { const assistant = assistantMessage("a1", 2); const tool = toolCall("tool-1", 3); @@ -302,13 +322,36 @@ describe("layoutStream", () => { }); expect(layout.auxiliaryTurnFooter).toBeNull(); - expect(findLayoutItem(layout, assistant.id).completedFooter?.itemId).toBe(assistant.id); - expect(footerOwners(layout)).toEqual([assistant.id]); + expect(findLayoutItem(layout, assistant.id).completedFooter).toBeNull(); + expect(findLayoutItem(layout, tool.id).completedFooter?.itemId).toBe(assistant.id); + expect(footerOwners(layout)).toEqual([tool.id]); + expect(footerAssistantIds(layout)).toEqual([assistant.id]); }, ); it.each(["web", "android"] as const)( - "uses the latest assistant for an inline footer when a turn has multiple assistant blocks on %s", + "places split live-head tool footer using the assistant from history on %s", + (platform) => { + const assistant = assistantMessage("a1", 2); + const tool = toolCall("tool-1", 3); + const layout = layoutFor({ + platform, + tail: [userMessage("u1", 1), assistant], + head: [tool, userMessage("u2", 4)], + timingIds: [assistant.id], + }); + + expect(layout.auxiliaryTurnFooter).toBeNull(); + expect(findLayoutItem(layout, assistant.id).completedFooter).toBeNull(); + expect(findLayoutItem(layout, tool.id).completedFooter?.itemId).toBe(assistant.id); + expect(inlineFooterPlacementByItemId(layout)).toEqual({ + [tool.id]: assistant.id, + }); + }, + ); + + it.each(["web", "android"] as const)( + "uses the latest assistant for footer content while placing after the visible turn end on %s", (platform) => { const firstAssistant = assistantMessage("a1", 2); const firstTool = toolCall("tool-1", 3); @@ -329,10 +372,46 @@ describe("layoutStream", () => { expect(layout.auxiliaryTurnFooter).toBeNull(); expect(findLayoutItem(layout, firstAssistant.id).completedFooter).toBeNull(); - expect(findLayoutItem(layout, latestAssistant.id).completedFooter?.itemId).toBe( + expect(findLayoutItem(layout, latestAssistant.id).completedFooter).toBeNull(); + expect(findLayoutItem(layout, latestTool.id).completedFooter?.itemId).toBe( latestAssistant.id, ); - expect(footerOwners(layout)).toEqual([latestAssistant.id]); + expect(footerOwners(layout)).toEqual([latestTool.id]); + expect(footerAssistantIds(layout)).toEqual([latestAssistant.id]); + }, + ); + + it.each(["web", "android"] as const)( + "keeps every completed turn footer while placing each one after that turn's last visible item on %s", + (platform) => { + const firstAssistant = assistantMessage("a1", 2); + const secondAssistant = assistantMessage("a2", 4); + const secondTool = toolCall("tool-2", 5); + const layout = layoutFor({ + platform, + tail: [ + userMessage("u1", 1), + firstAssistant, + userMessage("u2", 3), + secondAssistant, + secondTool, + userMessage("u3", 6), + ], + timingIds: [firstAssistant.id, secondAssistant.id], + }); + + expect(layout.auxiliaryTurnFooter).toBeNull(); + expect(findLayoutItem(layout, firstAssistant.id).completedFooter?.itemId).toBe( + firstAssistant.id, + ); + expect(findLayoutItem(layout, secondAssistant.id).completedFooter).toBeNull(); + expect(findLayoutItem(layout, secondTool.id).completedFooter?.itemId).toBe( + secondAssistant.id, + ); + expect(inlineFooterPlacementByItemId(layout)).toEqual({ + [firstAssistant.id]: firstAssistant.id, + [secondTool.id]: secondAssistant.id, + }); }, ); diff --git a/packages/app/src/agent-stream/layout.ts b/packages/app/src/agent-stream/layout.ts index 26ac70830..f276db382 100644 --- a/packages/app/src/agent-stream/layout.ts +++ b/packages/app/src/agent-stream/layout.ts @@ -44,7 +44,6 @@ export interface StreamLayoutInput { interface LayoutSegmentInput { strategy: StreamStrategy; - agentStatus: string; items: StreamItem[]; timingByAssistantId: Map; auxiliaryTurnFooter: TurnFooterHost | null; @@ -52,6 +51,14 @@ interface LayoutSegmentInput { boundaryIndex: number | null; boundaryAboveItem: StreamItem | null; boundaryBelowItem: StreamItem | null; + boundaryAboveItems: StreamItem[] | null; + boundaryAboveIndex: number | null; +} + +interface AssistantFooterSource { + item: Extract; + items: StreamItem[]; + index: number; } function createTurnFooterHost(input: { @@ -68,25 +75,45 @@ function createTurnFooterHost(input: { }; } -function findLatestAssistantIndexInTurn(input: { +function findLatestAssistantInTurn(input: { strategy: StreamStrategy; items: StreamItem[]; startIndex: number; -}): number | null { - for ( - let index = input.startIndex; - index >= 0 && index < input.items.length; - index = input.strategy.getNeighborIndex(index, "above") - ) { - const item = input.items[index]; - if (!item || item.kind === "user_message") { + boundaryAboveItems?: StreamItem[] | null; + boundaryAboveIndex?: number | null; +}): AssistantFooterSource | null { + let items = input.items; + let index = input.startIndex; + let canCrossBoundary = true; + + while (true) { + for ( + ; + index >= 0 && index < items.length; + index = input.strategy.getNeighborIndex(index, "above") + ) { + const item = items[index]; + if (!item || item.kind === "user_message") { + return null; + } + if (item.kind === "assistant_message") { + return { item, items, index }; + } + } + + if ( + !canCrossBoundary || + !input.boundaryAboveItems || + input.boundaryAboveIndex === null || + input.boundaryAboveIndex === undefined + ) { return null; } - if (item.kind === "assistant_message") { - return index; - } + + items = input.boundaryAboveItems; + index = input.boundaryAboveIndex; + canCrossBoundary = false; } - return null; } function resolveAuxiliaryTurnFooter(input: StreamLayoutInput): TurnFooterHost | null { @@ -100,104 +127,54 @@ function resolveAuxiliaryTurnFooter(input: StreamLayoutInput): TurnFooterHost | return null; } - const assistantIndex = findLatestAssistantIndexInTurn({ + const assistant = findLatestAssistantInTurn({ strategy: input.strategy, items: footerItems, startIndex: latestIndex, }); - if (assistantIndex === null) { - return null; - } - - const item = footerItems[assistantIndex]; - if (!item || item.kind !== "assistant_message") { + if (!assistant) { return null; } return createTurnFooterHost({ - item, - items: footerItems, - index: assistantIndex, + item: assistant.item, + items: assistant.items, + index: assistant.index, timingByAssistantId: input.timingByAssistantId, }); } -function findTurnEndIndexInSegment(input: { - strategy: StreamStrategy; - items: StreamItem[]; - startIndex: number; -}): number { - let endIndex = input.startIndex; - for ( - let index = input.strategy.getNeighborIndex(input.startIndex, "below"); - index >= 0 && index < input.items.length; - index = input.strategy.getNeighborIndex(index, "below") - ) { - const item = input.items[index]; - if (!item || item.kind === "user_message") { - break; - } - endIndex = index; - } - return endIndex; -} - -function shouldRenderCompletedFooter(input: { +function resolveCompletedFooter(input: { strategy: StreamStrategy; items: StreamItem[]; index: number; item: StreamItem; belowItem: StreamItem | null; - agentStatus: string; + timingByAssistantId: Map; auxiliaryTurnFooter: TurnFooterHost | null; - boundaryIndex: number | null; - boundaryBelowItem: StreamItem | null; -}): boolean { - if ( - input.item.kind !== "assistant_message" || - input.auxiliaryTurnFooter?.itemId === input.item.id - ) { - return false; + boundaryAboveItems: StreamItem[] | null; + boundaryAboveIndex: number | null; +}): TurnFooterHost | null { + if (input.item.kind === "user_message" || input.belowItem?.kind !== "user_message") { + return null; } - if ( - input.belowItem?.kind === "user_message" || - (input.belowItem === null && input.agentStatus !== "running") - ) { - return true; - } - - if (!isToolSequenceItem(input.belowItem)) { - return false; - } - - const sameSegmentBelowItem = input.strategy.getNeighborItem(input.items, input.index, "below"); - if (sameSegmentBelowItem?.id !== input.belowItem.id) { - return false; - } - - const turnEndIndex = findTurnEndIndexInSegment({ + const assistant = findLatestAssistantInTurn({ strategy: input.strategy, items: input.items, startIndex: input.index, + boundaryAboveItems: input.boundaryAboveItems, + boundaryAboveIndex: input.boundaryAboveIndex, }); - const belowTurnItem = getSegmentNeighbor({ - strategy: input.strategy, - items: input.items, - index: turnEndIndex, - relation: "below", - boundaryIndex: input.boundaryIndex, - boundaryItem: input.boundaryBelowItem, - }); - if (input.agentStatus === "running" && belowTurnItem?.kind !== "user_message") { - return false; + if (!assistant || input.auxiliaryTurnFooter?.itemId === assistant.item.id) { + return null; } - const assistantIndex = findLatestAssistantIndexInTurn({ - strategy: input.strategy, - items: input.items, - startIndex: turnEndIndex, + return createTurnFooterHost({ + item: assistant.item, + items: assistant.items, + index: assistant.index, + timingByAssistantId: input.timingByAssistantId, }); - return assistantIndex === input.index; } function isToolSequenceItem( @@ -270,24 +247,17 @@ function layoutSegment(input: LayoutSegmentInput): StreamLayoutItem[] { aboveItem, belowItem, }); - const completedFooter = shouldRenderCompletedFooter({ + const completedFooter = resolveCompletedFooter({ strategy: input.strategy, items: input.items, index, item, belowItem, - agentStatus: input.agentStatus, + timingByAssistantId: input.timingByAssistantId, auxiliaryTurnFooter: input.auxiliaryTurnFooter, - boundaryIndex: input.boundaryIndex, - boundaryBelowItem: input.boundaryBelowItem, - }) - ? createTurnFooterHost({ - item, - items: input.items, - index, - timingByAssistantId: input.timingByAssistantId, - }) - : null; + boundaryAboveItems: input.boundaryAboveItems, + boundaryAboveIndex: input.boundaryAboveIndex, + }); return { item, @@ -328,7 +298,6 @@ export function layoutStream(input: StreamLayoutInput): StreamLayout { // and .kind are stable across text-only flushes (text growth doesn't change what kind of // item borders history), so cached layout stays valid between flushes. const historyCacheKey = [ - input.agentStatus, frameOrder, historyBoundaryIndex ?? "null", liveHeadBoundaryItem?.id ?? "null", @@ -346,7 +315,6 @@ export function layoutStream(input: StreamLayoutInput): StreamLayout { } else { history = layoutSegment({ strategy: input.strategy, - agentStatus: input.agentStatus, items: input.history, timingByAssistantId: input.timingByAssistantId, auxiliaryTurnFooter, @@ -354,6 +322,8 @@ export function layoutStream(input: StreamLayoutInput): StreamLayout { boundaryIndex: historyBoundaryIndex, boundaryAboveItem: null, boundaryBelowItem: liveHeadBoundaryItem, + boundaryAboveItems: null, + boundaryAboveIndex: null, }); byKey.set(historyCacheKey, history); } @@ -363,7 +333,6 @@ export function layoutStream(input: StreamLayoutInput): StreamLayout { const liveHead = layoutSegment({ strategy: input.strategy, - agentStatus: input.agentStatus, items: input.liveHead, timingByAssistantId: input.timingByAssistantId, auxiliaryTurnFooter, @@ -371,6 +340,8 @@ export function layoutStream(input: StreamLayoutInput): StreamLayout { boundaryIndex: liveHeadBoundaryIndex, boundaryAboveItem: historyBoundaryItem, boundaryBelowItem: null, + boundaryAboveItems: input.history, + boundaryAboveIndex: historyBoundaryIndex, }); return {