@truefoundry/assistant-ui-runtime 0.1.15 → 0.1.16

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@truefoundry/assistant-ui-runtime",
-  "version": "0.1.15",
+  "version": "0.1.16",
   "description": "TrueFoundry Gateway agent runtime adapter for assistant-ui",
   "license": "Apache-2.0",
   "repository": {

package/src/convertTurnMessages.test.ts

@@ -18,8 +18,8 @@ import type {
 import { ROOT_THREAD_ID } from "./constants.js";
 import { collectPendingToolResponses } from "./collectPending.js";
 import {
-    buildSnapshotBeforeTurnIndex,
     buildSnapshotFromSessionEvents,
+    buildSnapshotThroughTurn,
     buildTurnAssistantContent,
     buildUserMessageContent,
     buildUserMessageFromTurnInput,
@@ -28,6 +28,7 @@ import {
     prependOlderSessionHistory,
     projectSessionMessages,
     repositoryItemsFromMessages,
+    resolveGatewayBranchPreviousTurnIdForTurn,
     streamTurnEvents,
     turnStreamUpdateToAssistantMessage,
 } from "./convertTurnMessages.js";
@@ -135,22 +136,68 @@ function mockTurn(
     };
 }
 
+/** Mirror runtime listTurns normalization: ASC stays, DESC (or tied) reverses. */
+function toOldestFirstTurns(turns: TurnFixture[]): TurnFixture[] {
+    const chain = turns.filter((turn) => turn.state.status !== "running");
+    let newestFirst = chain.length > 1;
+    for (let i = 1; i < chain.length; i++) {
+        const newer = chain[i];
+        const older = chain[i - 1];
+        if (newer == null || older == null) {
+            continue;
+        }
+        const delta = Date.parse(newer.createdAt) - Date.parse(older.createdAt);
+        if (delta < 0) {
+            newestFirst = true;
+            break;
+        }
+        if (delta > 0) {
+            newestFirst = false;
+            break;
+        }
+    }
+    return newestFirst ? [...chain].reverse() : [...chain];
+}
+
+/** Mirrors the host's ancestor walk: parent pointers only, sibling branches excluded. */
+function ancestorChainFromTurns(
+    turns: TurnFixture[],
+    anchorTurnId: string,
+): TurnFixture[] {
+    const turnsById = new Map(turns.map((turn) => [turn.id, turn]));
+    const chain: TurnFixture[] = [];
+    const seen = new Set<string>();
+    let cursor: string | null | undefined = anchorTurnId;
+    while (cursor != null && !seen.has(cursor)) {
+        seen.add(cursor);
+        const turn = turnsById.get(cursor);
+        if (turn == null) {
+            break;
+        }
+        if (turn.state.status !== "running") {
+            chain.unshift(turn);
+        }
+        cursor = turn.previousTurnId;
+    }
+    return chain;
+}
+
 /**
- * Builds session-level event items from per-turn mocks — newest-first turns,
- * running turns excluded (matches session.listEvents API contract).
+ * Builds session-level event items from per-turn mocks — accepts ASC or DESC
+ * listTurns order; running turns excluded (matches session.listEvents). An
+ * anchor selects that turn's ancestor chain, not a chronological prefix.
  */
 function sessionEventItemsFromTurns(
-    turnsNewestFirst: TurnFixture[],
+    turns: TurnFixture[],
     lastTurnId?: string,
 ): SessionEventItem[] {
-    let chain = turnsNewestFirst.filter((turn) => turn.state.status !== "running");
-    if (lastTurnId != null) {
-        const anchorIndex = chain.findIndex((turn) => turn.id === lastTurnId);
-        chain = anchorIndex === -1 ? [] : chain.slice(anchorIndex);
-    }
+    const chain =
+        lastTurnId == null
+            ? toOldestFirstTurns(turns)
+            : ancestorChainFromTurns(turns, lastTurnId);
 
     const items: SessionEventItem[] = [];
-    for (const turn of [...chain].reverse()) {
+    for (const turn of chain) {
         items.push({
             turnId: turn.id,
             event: {
@@ -160,6 +207,9 @@ function sessionEventItemsFromTurns(
                 input: turn.input,
                 state: { status: "running" },
                 createdAt: turn.createdAt,
+                ...(turn.previousTurnId === undefined
+                    ? {}
+                    : { previousTurnId: turn.previousTurnId }),
             },
         });
         for (const event of turn.events) {
@@ -180,10 +230,18 @@ function sessionEventItemsFromTurns(
 
 function mockServerWithTurns(turns: TurnFixture[]): AgentChatServer {
     const eventsByTurnId = new Map(turns.map((turn) => [turn.id, turn.events]));
+    const turnsById = new Map(turns.map((turn) => [turn.id, turn]));
     return {
         listTurns: async ({ limit }: { limit?: number } = {}) => ({
             data: limit != null ? turns.slice(0, limit) : turns,
         }),
+        getTurn: async ({ turnId }: { turnId: string }) => {
+            const turn = turnsById.get(turnId);
+            if (turn == null) {
+                throw new Error(`Turn ${turnId} not found`);
+            }
+            return turn;
+        },
         listEvents: async (opts: { lastTurnId?: string } = {}) => {
             const items = sessionEventItemsFromTurns(turns, opts.lastTurnId);
             const newestFirst = [...items].reverse();
@@ -2824,40 +2882,46 @@ describe("buildSnapshotFromSessionEvents", () => {
     });
 });
 
-describe("buildSnapshotBeforeTurnIndex", () => {
-    it("rewinds via session.listEvents({ lastTurnId }) excluding the branch turn", async () => {
-        const t1 = mockTurn({
-            id: "t1",
-            createdAt,
-            input: [{ type: "user.message", content: "first" }],
-            events: [
-                modelMessage({ id: "m1", threadId: ROOT_THREAD_ID, content: "reply 1" }),
+describe("buildSnapshotThroughTurn", () => {
+    /** Linear session t1 -> t2 -> t3, oldest-first. */
+    function linearTurns(): TurnFixture[] {
+        return [
+            mockTurn({
+                id: "t1",
+                createdAt: "2026-01-01T00:00:00.000Z",
+                previousTurnId: null,
+                input: [{ type: "user.message", content: "first" }],
+                events: [
+                    modelMessage({ id: "m1", threadId: ROOT_THREAD_ID, content: "reply 1" }),
                 ],
-        });
-        const t2 = mockTurn({
-            id: "t2",
-            createdAt,
-            input: [{ type: "user.message", content: "second" }],
-            events: [
-                modelMessage({ id: "m2", threadId: ROOT_THREAD_ID, content: "reply 2" }),
+            }),
+            mockTurn({
+                id: "t2",
+                createdAt: "2026-01-01T00:01:00.000Z",
+                previousTurnId: "t1",
+                input: [{ type: "user.message", content: "second" }],
+                events: [
+                    modelMessage({ id: "m2", threadId: ROOT_THREAD_ID, content: "reply 2" }),
                 ],
-        });
-        const t3 = mockTurn({
-            id: "t3",
-            createdAt,
-            input: [{ type: "user.message", content: "third" }],
-            events: [
-                modelMessage({ id: "m3", threadId: ROOT_THREAD_ID, content: "reply 3" }),
+            }),
+            mockTurn({
+                id: "t3",
+                createdAt: "2026-01-01T00:02:00.000Z",
+                previousTurnId: "t2",
+                input: [{ type: "user.message", content: "third" }],
+                events: [
+                    modelMessage({ id: "m3", threadId: ROOT_THREAD_ID, content: "reply 3" }),
                 ],
-        });
+            }),
+        ];
+    }
 
-        // listTurns is newest-first.
-        const server = mockServerWithTurns([t3, t2, t1]);
-        const snapshot = await buildSnapshotBeforeTurnIndex(server, SESSION_ID, 2);
+    it("rewinds to the anchor's ancestor chain, inclusive", async () => {
+        const server = mockServerWithTurns(linearTurns());
+        const snapshot = await buildSnapshotThroughTurn(server, SESSION_ID, "t2");
 
         expect(snapshot.turns.map((turn) => turn.id)).toEqual(["t1", "t2"]);
-        const messages = projectSessionMessages(snapshot);
-        expect(messages.map((message) => message.id)).toEqual([
+        expect(projectSessionMessages(snapshot).map((message) => message.id)).toEqual([
             "t1-user",
             "t1-assistant",
             "t2-user",
@@ -2865,12 +2929,116 @@ describe("buildSnapshotBeforeTurnIndex", () => {
         ]);
     });
 
-    it("returns an empty snapshot when branching from the first turn", async () => {
-        const server = mockServerWithTurns([
-            mockTurn({ id: "t1", createdAt }),
-        ]);
-        const snapshot = await buildSnapshotBeforeTurnIndex(server, SESSION_ID, 0);
+    it("returns an empty snapshot for a null anchor (branching from a root turn)", async () => {
+        const server = mockServerWithTurns(linearTurns());
+        const snapshot = await buildSnapshotThroughTurn(server, SESSION_ID, null);
         expect(snapshot.turns).toHaveLength(0);
     });
+
+    it("ignores listTurns page order", async () => {
+        const asc = linearTurns();
+        const ascSnapshot = await buildSnapshotThroughTurn(
+            mockServerWithTurns(asc),
+            SESSION_ID,
+            "t2",
+        );
+        const descSnapshot = await buildSnapshotThroughTurn(
+            mockServerWithTurns([...asc].reverse()),
+            SESSION_ID,
+            "t2",
+        );
+
+        expect(descSnapshot.turns.map((turn) => turn.id)).toEqual(
+            ascSnapshot.turns.map((turn) => turn.id),
+        );
+        expect(descSnapshot.turns.map((turn) => turn.id)).toEqual(["t1", "t2"]);
+    });
+
+    it("drops sibling turns left behind by an earlier rerun", async () => {
+        // m2/m3 were rerun once already: t2/t3 are orphans, t2p/t3p are live.
+        const [t1, t2, t3] = linearTurns();
+        const t2p = mockTurn({
+            id: "t2p",
+            createdAt: "2026-01-01T00:03:00.000Z",
+            previousTurnId: "t1",
+            input: [{ type: "user.message", content: "second again" }],
+            events: [
+                modelMessage({ id: "m2p", threadId: ROOT_THREAD_ID, content: "reply 2p" }),
+            ],
+        });
+        const t3p = mockTurn({
+            id: "t3p",
+            createdAt: "2026-01-01T00:04:00.000Z",
+            previousTurnId: "t2p",
+            input: [{ type: "user.message", content: "third again" }],
+            events: [
+                modelMessage({ id: "m3p", threadId: ROOT_THREAD_ID, content: "reply 3p" }),
+            ],
+        });
+        const server = mockServerWithTurns([t1!, t2!, t3!, t2p, t3p]);
+
+        // Rerunning the middle live turn rewinds to its parent, not to the
+        // chronologically preceding turn (t3), whose chain still holds t2/t3.
+        const previousTurnId = await resolveGatewayBranchPreviousTurnIdForTurn(
+            server,
+            SESSION_ID,
+            "t2p",
+        );
+        expect(previousTurnId).toBe("t1");
+
+        const snapshot = await buildSnapshotThroughTurn(server, SESSION_ID, previousTurnId);
+        expect(snapshot.turns.map((turn) => turn.id)).toEqual(["t1"]);
+
+        // The orphan chain is still reachable, so the assertion above is not
+        // passing by accident.
+        const orphanSnapshot = await buildSnapshotThroughTurn(server, SESSION_ID, "t3");
+        expect(orphanSnapshot.turns.map((turn) => turn.id)).toEqual(["t1", "t2", "t3"]);
+    });
+});
+
+describe("resolveGatewayBranchPreviousTurnIdForTurn", () => {
+    it("uses the turn parent for subsequent turns and none for roots (ASC listTurns)", async () => {
+        const t1 = mockTurn({
+            id: "t1",
+            createdAt: "2026-01-01T00:00:00.000Z",
+            previousTurnId: null,
+        });
+        const t2 = mockTurn({
+            id: "t2",
+            createdAt: "2026-01-01T00:01:00.000Z",
+            previousTurnId: "t1",
+        });
+        const t3 = mockTurn({
+            id: "t3",
+            createdAt: "2026-01-01T00:02:00.000Z",
+            previousTurnId: "t2",
+        });
+        const server = mockServerWithTurns([t1, t2, t3]);
+
+        await expect(
+            resolveGatewayBranchPreviousTurnIdForTurn(server, SESSION_ID, "t1"),
+        ).resolves.toBe("none");
+        await expect(
+            resolveGatewayBranchPreviousTurnIdForTurn(server, SESSION_ID, "t3"),
+        ).resolves.toBe("t2");
+    });
+
+    it("uses the turn parent when listTurns is newest-first", async () => {
+        const t1 = mockTurn({
+            id: "t1",
+            createdAt: "2026-01-01T00:00:00.000Z",
+            previousTurnId: null,
+        });
+        const t2 = mockTurn({
+            id: "t2",
+            createdAt: "2026-01-01T00:01:00.000Z",
+            previousTurnId: "t1",
+        });
+        const server = mockServerWithTurns([t2, t1]);
+
+        await expect(
+            resolveGatewayBranchPreviousTurnIdForTurn(server, SESSION_ID, "t2"),
+        ).resolves.toBe("t1");
+    });
 });
 

package/src/convertTurnMessages.ts

@@ -1379,93 +1379,38 @@ export async function buildSnapshotFromSession(
     });
 }
 
-/** Rebuilds session state from turns strictly before `beforeTurnId` (excludes that turn). */
-export async function buildSnapshotBeforeTurn(
-    server: AgentChatServer,
-    sessionId: string,
-    beforeTurnId: string,
-    concurrency: number = DEFAULT_LIST_EVENTS_CONCURRENCY,
-): Promise<SessionSnapshot> {
-    const turns = await listSessionTurnsOrdered(server, sessionId);
-
-    const beforeIndex = turns.findIndex((turn) => turn.id === beforeTurnId);
-    if (beforeIndex === -1) {
-        throw new Error(`Turn ${beforeTurnId} not found in session`);
-    }
-
-    return buildSnapshotBeforeTurnIndex(
-        server,
-        sessionId,
-        beforeIndex,
-        concurrency,
-        turns,
-    );
-}
-
-/** Rebuilds session state from the first `turnIndex` turns (excludes that turn). */
-export async function buildSnapshotBeforeTurnIndex(
+/**
+ * Rebuilds the conversation through `anchorTurnId`, including that turn.
+ * The server follows parent links from the anchor, so turns from abandoned
+ * branches are excluded. A null anchor represents an empty conversation.
+ */
+export async function buildSnapshotThroughTurn(
     server: AgentChatServer,
     sessionId: string,
-    turnIndex: number,
-    _concurrency: number = DEFAULT_LIST_EVENTS_CONCURRENCY,
-    orderedTurns?: Turn[],
+    anchorTurnId: string | null,
 ): Promise<SessionSnapshot> {
-    if (turnIndex <= 0) {
+    if (anchorTurnId == null) {
         return createEmptySessionSnapshot();
     }
-
-    const turns = orderedTurns ?? (await listSessionTurnsOrdered(server, sessionId));
-    const turnsToInclude = turns.slice(0, turnIndex);
-    const lastTurnId = turnsToInclude.at(-1)?.id;
-    if (lastTurnId == null) {
-        return createEmptySessionSnapshot();
-    }
-
-    // Anchor the session events window at the newest included turn so the
-    // ancestor chain matches `[turns[0], …, turns[turnIndex - 1]]`.
-    const items = await fetchAllSessionEvents(server, sessionId, { lastTurnId });
+    const items = await fetchAllSessionEvents(server, sessionId, {
+        lastTurnId: anchorTurnId,
+    });
     const snapshot = createEmptySessionSnapshot();
     ingestSessionEventsIntoSnapshot(snapshot, items);
     return snapshot;
 }
 
-/** Turn id to branch from when resubmitting at `turnIndex` (`"none"` for first turn). */
-export async function resolveGatewayBranchPreviousTurnId(
-    server: AgentChatServer,
-    sessionId: string,
-    turnIndex: number,
-    orderedTurns?: Turn[],
-): Promise<string> {
-    if (turnIndex <= 0) {
-        return "none";
-    }
-    const turns = orderedTurns ?? (await listSessionTurnsOrdered(server, sessionId));
-    return turns[turnIndex - 1]?.id ?? "none";
-}
-
-/** Resolves `previousTurnId` by turn id so partial history windows stay correct. */
+/**
+ * Resolves `previousTurnId` for edit/retry of `turnId` from the turn's own
+ * parent pointer (`"none"` for roots). Independent of listTurns order.
+ */
 export async function resolveGatewayBranchPreviousTurnIdForTurn(
     server: AgentChatServer,
     sessionId: string,
     turnId: string,
 ): Promise<string> {
-    const turns = await listSessionTurnsOrdered(server, sessionId);
-    const turnIndex = turns.findIndex((turn) => turn.id === turnId);
-    return resolveGatewayBranchPreviousTurnId(server, sessionId, turnIndex, turns);
-}
-
-async function listSessionTurnsOrdered(
-    server: AgentChatServer,
-    sessionId: string,
-): Promise<Turn[]> {
-    const turns = await drainListPages((pageToken) =>
-        server.listTurns({
-            sessionId,
-            ...(pageToken != null ? { pageToken } : {}),
-        }),
-    );
-    turns.reverse();
-    return turns;
+    const turn = await server.getTurn({ sessionId, turnId });
+    return turn.previousTurnId ?? "none";
 }
 
 export async function buildTurnAssistantContent(

package/src/types.ts

@@ -28,7 +28,6 @@ type TrueFoundryAgentRuntimeBaseOptions = ExternalStoreSharedOptions & {
     threadId?: string | undefined;
     onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
     onError?: ((error: unknown) => void) | undefined;
-    listEventsConcurrency?: number | undefined;
     /**
      * Optional filter forwarded to `listSessions({ agentId })`.
      * Omit for all chats; hosts that key agents by name pass that name as the id.

package/src/useTrueFoundryAgentMessages.test.tsx

@@ -47,6 +47,7 @@ vi.mock("./convertTurnMessages.js", async (importOriginal) => {
 const mockServer = {
     cancelSession: vi.fn().mockResolvedValue(undefined),
     listTurns: vi.fn(),
+    getTurn: vi.fn(),
     // Present so resume-capable paths are exercised; resumeTurnStream is mocked.
     subscribeToTurn: vi.fn(),
 } as unknown as AgentChatServer;
@@ -429,6 +430,7 @@ describe("useTrueFoundryAgentMessages", () => {
             expect.any(PeerThreadFoldState),
             {
                 userMessage: "first",
+                previousTurnId: "none",
                 headers: {
                     "x-tfy-session-last-updated-at": "2026-06-30T12:00:00.000Z",
                 },
@@ -689,21 +691,22 @@ describe("useTrueFoundryAgentMessages", () => {
 
     it("editFromTurn drops prior turns before showing the edited user message", async () => {
         const createdAt = new Date().toISOString();
+        const rootTurn = {
+            id: "turn-1",
+            sessionId: "session-1",
+            createdAt,
+            previousTurnId: null,
+            state: {
+                status: "done" as const,
+                requiredActions: [],
+                completedAt: createdAt,
+            },
+            input: [{ type: "user.message" as const, content: "Hello" }],
+        } as Turn;
         vi.mocked(mockServer.listTurns).mockResolvedValue({
-            data: [
-                {
-                    id: "turn-1",
-                    sessionId: "session-1",
-                    createdAt,
-                    state: {
-                        status: "done",
-                        requiredActions: [],
-                        completedAt: createdAt,
-                    },
-                    input: [{ type: "user.message", content: "Hello" }],
-                } as Turn,
-            ],
+            data: [rootTurn],
         });
+        vi.mocked(mockServer.getTurn).mockResolvedValue(rootTurn);
         const fold = new PeerThreadFoldState();
         ingestTurnEvent(fold, {
             type: "model.message",
@@ -805,21 +808,22 @@ describe("useTrueFoundryAgentMessages", () => {
         const onError = vi.fn();
         const original = snapshotWithUserTurn("Hello");
         vi.mocked(loadSessionSnapshot).mockResolvedValue(original);
+        const rootTurn = {
+            id: "turn-1",
+            sessionId: "session-1",
+            createdAt,
+            previousTurnId: null,
+            state: {
+                status: "done" as const,
+                requiredActions: [],
+                completedAt: createdAt,
+            },
+            input: [{ type: "user.message" as const, content: "Hello" }],
+        } as Turn;
         vi.mocked(mockServer.listTurns).mockResolvedValue({
-            data: [
-                {
-                    id: "turn-1",
-                    sessionId: "session-1",
-                    createdAt,
-                    state: {
-                        status: "done",
-                        requiredActions: [],
-                        completedAt: createdAt,
-                    },
-                    input: [{ type: "user.message", content: "Hello" }],
-                } as Turn,
-            ],
+            data: [rootTurn],
         });
+        vi.mocked(mockServer.getTurn).mockResolvedValue(rootTurn);
         vi.mocked(streamTurnContent).mockImplementation(async function* () {
             throw new Error("Turn preparation failed");
         });

package/src/useTrueFoundryAgentMessages.ts

@@ -15,7 +15,7 @@ import type { AgentChatServer } from "./server/types.js";
 import { ROOT_THREAD_ID } from "./constants.js";
 import {
     buildEditedUserMessageContent,
-    buildSnapshotBeforeTurn,
+    buildSnapshotThroughTurn,
     computeGroupRootBaseline,
     extractTurnUserMessageContent,
     prependOlderSessionHistory,
@@ -60,7 +60,6 @@ export type UseTrueFoundryAgentMessagesOptions = {
     isMain?: boolean | undefined;
     /** URL-selected session may load before the thread list marks it as main. */
     isInitialSession?: boolean | undefined;
-    listEventsConcurrency?: number | undefined;
     onError?: ((error: unknown) => void) | undefined;
     initializeSession?: () => Promise<{
         remoteId: string;
@@ -293,7 +292,6 @@ export function useTrueFoundryAgentMessages({
     sessionId,
     isMain,
     isInitialSession,
-    listEventsConcurrency,
     onError,
     initializeSession,
     resolveConversationSessionId,
@@ -975,11 +973,12 @@ export function useTrueFoundryAgentMessages({
                     conversationSessionId,
                     turnId,
                 );
-                rewound = await buildSnapshotBeforeTurn(
+                // Rewind to the exact parent used for the new branch. Using the
+                // previous item from listTurns could select an abandoned branch.
+                rewound = await buildSnapshotThroughTurn(
                     server,
                     conversationSessionId,
-                    turnId,
-                    listEventsConcurrency,
+                    previousTurnId === "none" ? null : previousTurnId,
                 );
                 createdAtByMessageIdRef.current = new Map();
                 // Keep the ref aligned before awaiting sendTurn so any intermediate
@@ -1000,13 +999,7 @@ export function useTrueFoundryAgentMessages({
                 branchRollbackSnapshot: committed,
             });
         },
-        [
-            cancel,
-            server,
-            listEventsConcurrency,
-            sendTurn,
-            sessionId,
-        ],
+        [cancel, server, sendTurn, sessionId],
     );
 
     const resetFromTurn = useCallback(

package/src/useTrueFoundryAgentRuntime.ts

@@ -49,7 +49,6 @@ function useTrueFoundryAgentRuntimeImpl(
         agent,
         adapters,
         onError,
-        listEventsConcurrency,
         ...sharedOptions
     } = options;
 
@@ -126,7 +125,6 @@ function useTrueFoundryAgentRuntimeImpl(
         sessionId,
         isMain,
         isInitialSession,
-        listEventsConcurrency,
         onError,
         initializeSession,
         getTurnHeaders: agent.mode === "draft" ? getTurnHeaders : undefined,