fix: stabilize chat turn projection

2026-05-07 12:10:04 +03:00
parent 4cc2f1bff6
commit 07edd6e9b9
3 changed files with 168 additions and 42 deletions
@@ -0,0 +1,89 @@
+import { describe, expect, test } from 'bun:test';
+import type { Message, Part } from '@opencode-ai/sdk/v2';
+import { projectTurnRecords } from './projectTurnRecords';
+import type { ChatMessageEntry } from './types';
+
+function createMessageEntry({
+    id,
+    role,
+    parentID,
+    createdAt,
+}: {
+    id: string;
+    role: 'user' | 'assistant' | 'system';
+    parentID?: string;
+    createdAt: number;
+}): ChatMessageEntry {
+    return {
+        info: {
+            id,
+            role,
+            ...(parentID ? { parentID } : {}),
+            time: { created: createdAt },
+        } as Message,
+        parts: [] as Part[],
+    };
+}
+
+describe('projectTurnRecords', () => {
+    test('groups assistant replies under their parent user turn', () => {
+        const user = createMessageEntry({ id: 'u1', role: 'user', createdAt: 1 });
+        const assistant = createMessageEntry({ id: 'a1', role: 'assistant', parentID: 'u1', createdAt: 2 });
+
+        const projection = projectTurnRecords([user, assistant]);
+
+        expect(projection.turns).toHaveLength(1);
+        expect(projection.turns[0]?.turnId).toBe('u1');
+        expect(projection.turns[0]?.assistantMessageIds).toEqual(['a1']);
+        expect(projection.ungroupedMessageIds.size).toBe(0);
+    });
+
+    test('keeps out-of-order assistant replies attached to their parent user turn', () => {
+        const user1 = createMessageEntry({ id: 'u1', role: 'user', createdAt: 1 });
+        const assistant1 = createMessageEntry({ id: 'a1', role: 'assistant', parentID: 'u1', createdAt: 2 });
+        const assistant2 = createMessageEntry({ id: 'a2', role: 'assistant', parentID: 'u2', createdAt: 4 });
+        const user2 = createMessageEntry({ id: 'u2', role: 'user', createdAt: 3 });
+
+        const projection = projectTurnRecords([user1, assistant1, assistant2, user2]);
+
+        expect(projection.turns).toHaveLength(2);
+        expect(projection.turns[0]?.turnId).toBe('u1');
+        expect(projection.turns[0]?.assistantMessageIds).toEqual(['a1']);
+        expect(projection.turns[1]?.turnId).toBe('u2');
+        expect(projection.turns[1]?.assistantMessageIds).toEqual(['a2']);
+        expect(projection.ungroupedMessageIds.size).toBe(0);
+    });
+
+    test('does not render assistant replies while their parent user turn is missing', () => {
+        const user1 = createMessageEntry({ id: 'u1', role: 'user', createdAt: 1 });
+        const assistant1 = createMessageEntry({ id: 'a1', role: 'assistant', parentID: 'u1', createdAt: 2 });
+        const assistant2 = createMessageEntry({ id: 'a2', role: 'assistant', parentID: 'u2', createdAt: 4 });
+
+        const projection = projectTurnRecords([user1, assistant1, assistant2]);
+
+        expect(projection.turns).toHaveLength(1);
+        expect(projection.turns[0]?.turnId).toBe('u1');
+        expect(projection.turns[0]?.assistantMessageIds).toEqual(['a1']);
+        expect(projection.ungroupedMessageIds.has('a2')).toBe(false);
+        expect(projection.indexes.messageToTurnId.has('a2')).toBe(false);
+    });
+
+    test('does not render orphan assistant messages as standalone ungrouped entries', () => {
+        const assistant = createMessageEntry({ id: 'a1', role: 'assistant', parentID: 'missing-user', createdAt: 1 });
+
+        const projection = projectTurnRecords([assistant]);
+
+        expect(projection.turns).toHaveLength(0);
+        expect(projection.ungroupedMessageIds.has('a1')).toBe(false);
+        expect(projection.indexes.messageToTurnId.has('a1')).toBe(false);
+    });
+
+    test('keeps non-assistant orphan messages available as ungrouped entries', () => {
+        const system = createMessageEntry({ id: 's1', role: 'system', createdAt: 1 });
+
+        const projection = projectTurnRecords([system]);
+
+        expect(projection.turns).toHaveLength(0);
+        expect(projection.ungroupedMessageIds.has('s1')).toBe(true);
+    });
+});
@@ -102,46 +102,47 @@ export const projectTurnRecords = (
    const turns: TurnRecord[] = [];
    const turnByUserId = new Map<string, TurnRecord>();
    const groupedMessageIds = new Set<string>();
-    let currentTurn: TurnRecord | undefined;

    messages.forEach((message, index) => {
        const role = resolveMessageRole(message);
-        if (role === 'user') {
-            const turnId = message.info.id;
-            const turn: TurnRecord = {
-                turnId,
-                userMessageId: message.info.id,
-                userMessage: message,
-                headerMessageId: undefined,
-                messages: [createTurnMessageRecord(message, index)],
-                assistantMessageIds: [],
-                assistantMessages: [],
-                activityParts: [],
-                activitySegments: [],
-                summary: {},
-                summaryText: undefined,
-                hasTools: false,
-                hasReasoning: false,
-                diffStats: undefined,
-                stream: {
-                    isStreaming: false,
-                    isRetrying: false,
-                },
-            };
-            turns.push(turn);
-            turnByUserId.set(turn.userMessageId, turn);
-            groupedMessageIds.add(message.info.id);
-            currentTurn = turn;
+        if (role !== 'user') {
            return;
        }

+        const turnId = message.info.id;
+        const turn: TurnRecord = {
+            turnId,
+            userMessageId: message.info.id,
+            userMessage: message,
+            headerMessageId: undefined,
+            messages: [createTurnMessageRecord(message, index)],
+            assistantMessageIds: [],
+            assistantMessages: [],
+            activityParts: [],
+            activitySegments: [],
+            summary: {},
+            summaryText: undefined,
+            hasTools: false,
+            hasReasoning: false,
+            diffStats: undefined,
+            stream: {
+                isStreaming: false,
+                isRetrying: false,
+            },
+        };
+        turns.push(turn);
+        turnByUserId.set(turn.userMessageId, turn);
+        groupedMessageIds.add(message.info.id);
+    });
+
+    messages.forEach((message, index) => {
+        const role = resolveMessageRole(message);
        if (role !== 'assistant') {
            return;
        }

        const parentId = getMessageParentId(message);
-        const parentTurn = parentId ? turnByUserId.get(parentId) : undefined;
-        const targetTurn = parentTurn ?? currentTurn;
+        const targetTurn = parentId ? turnByUserId.get(parentId) : undefined;
        if (!targetTurn) {
            return;
        }
@@ -153,10 +154,6 @@ export const projectTurnRecords = (
            targetTurn.headerMessageId = message.info.id;
        }
        groupedMessageIds.add(message.info.id);
-
-        if (!parentTurn) {
-            currentTurn = targetTurn;
-        }
    });

    turns.forEach((turn) => {
@@ -185,6 +182,9 @@ export const projectTurnRecords = (
    const projection = projectTurnIndexes(turns);
    const ungroupedMessageIds = new Set<string>();
    messages.forEach((message) => {
+        if (resolveMessageRole(message) === 'assistant') {
+            return;
+        }
        if (!groupedMessageIds.has(message.info.id)) {
            ungroupedMessageIds.add(message.info.id);
        }