@truefoundry/assistant-ui-runtime 0.1.14 → 0.1.16

package/src/convertTurnMessages.ts

@@ -516,10 +516,6 @@ function attachRunningTurn(
     if (runningTurn == null) {
         return snapshot;
     }
-    // Session-level history excludes the running turn. Apply continuation inputs
-    // from the turn listing so answered approvals / ask-user prompts are not
-    // restored as pending while reconnecting to that turn after a refresh.
-    applyUserToolResponsesToFold(snapshot.fold, runningTurn.input ?? []);
     const pendingUserText = extractTurnUserText(runningTurn.input);
     return replaceSessionSnapshot(snapshot, {
         runningTurn,
@@ -537,13 +533,117 @@ function attachRunningTurn(
     });
 }
 
+/**
+ * Last `turn.created` in ASC event order with no following `turn.done` — the
+ * open tip of the active branch in this window.
+ */
+function findOpenTurnCreated(
+    itemsAsc: readonly GatewaySessionEventItem[],
+): { turnId: string; event: TurnCreatedEvent } | undefined {
+    let open: { turnId: string; event: TurnCreatedEvent } | undefined;
+    for (const item of itemsAsc) {
+        if (item.event.type === "turn.created") {
+            open = { turnId: item.turnId, event: item.event };
+        } else if (item.event.type === "turn.done") {
+            open = undefined;
+        }
+    }
+    return open;
+}
+
+function turnFromCreatedEvent(options: {
+    sessionId: string;
+    turnId: string;
+    event: TurnCreatedEvent;
+}): Turn {
+    const { sessionId, turnId, event } = options;
+    return {
+        id: turnId,
+        sessionId,
+        state: { status: "running" },
+        createdAt: event.createdAt,
+        ...(event.input != null ? { input: event.input } : {}),
+        ...(event.previousTurnId === undefined
+            ? {}
+            : { previousTurnId: event.previousTurnId }),
+    };
+}
+
+type SessionTip = {
+    /** Turn to resume and subscribe to; absent once the tip has finished. */
+    runningTurn?: Turn;
+    /**
+     * Tip input that event ingestion could not apply, because it only folds a
+     * turn's input once that turn's `turn.done` arrives. Answered ask-user
+     * prompts and approvals live here, so this must be folded even when the tip
+     * is no longer running.
+     */
+    continuationInput: readonly TurnInputItem[];
+};
+
+/**
+ * Resolves the tip turn of the active branch for resume/subscribe.
+ *
+ * Prefer an open tip from the events window (works when listTurns is
+ * oldest-first and when listEvents includes the running turn). Fall back to
+ * `listTurns({ limit: 1 })` for hosts that omit the running turn from
+ * listEvents and put the tip first.
+ */
+async function resolveSessionTip(options: {
+    server: AgentChatServer;
+    sessionId: string;
+    itemsAsc: readonly GatewaySessionEventItem[];
+}): Promise<SessionTip> {
+    const { server, sessionId, itemsAsc } = options;
+    const open = findOpenTurnCreated(itemsAsc);
+    if (open != null) {
+        const continuationInput = open.event.input ?? [];
+        if (typeof server.getTurn === "function") {
+            try {
+                const turn = await server.getTurn({
+                    sessionId,
+                    turnId: open.turnId,
+                });
+                if (turn.state.status === "running") {
+                    return {
+                        runningTurn: turn,
+                        continuationInput: turn.input ?? continuationInput,
+                    };
+                }
+                // Tip finished between listEvents and getTurn: nothing to
+                // resume, but its answers still belong in the fold.
+                return { continuationInput: turn.input ?? continuationInput };
+            } catch {
+                // getTurn failed; synthesize from the open turn.created below.
+            }
+        }
+        return {
+            runningTurn: turnFromCreatedEvent({
+                sessionId,
+                turnId: open.turnId,
+                event: open.event,
+            }),
+            continuationInput,
+        };
+    }
+
+    // Hosts that exclude the running turn from listEvents still surface it as
+    // the first row of listTurns when that API is tip-first.
+    const turnsPage = await server.listTurns({ sessionId, limit: 1 });
+    const tip = turnsPage.data[0] as Turn | undefined;
+    return tip?.state?.status === "running"
+        ? { runningTurn: tip, continuationInput: tip.input ?? [] }
+        : { continuationInput: [] };
+}
+
 /**
  * Builds a session snapshot using the session-level `listEvents` API.
  *
  * Loads the newest event page (extending only when a page boundary splits a
  * turn group), then leaves older pages for `prependOlderSessionHistory`.
- * Only `listTurns({ limit: 1 })` is called first to detect a currently-running
- * turn (the session-level API does not return events for the running turn).
+ * Detects a currently-running tip from an open `turn.created` in that window
+ * when present; otherwise falls back to `listTurns({ limit: 1 })` for hosts
+ * that omit the running turn from listEvents.
  *
  * `onProgress` is called after each complete turn is ingested so callers can
  * update the UI progressively while the processing loop runs.
@@ -553,12 +653,6 @@ export async function buildSnapshotFromSessionEvents(
     sessionId: string,
     onProgress?: (snap: SessionSnapshot) => void,
 ): Promise<SessionSnapshot> {
-    // Detect a running turn with a single listTurns page — do not drain pagination.
-    const turnsPage = await server.listTurns({ sessionId, limit: 1 });
-    const newestTurn = turnsPage.data[0] as Turn | undefined;
-    const runningTurn =
-        newestTurn?.state?.status === "running" ? newestTurn : undefined;
-
     const window = await fetchSessionEventsWindow(server, sessionId);
     const historyPagination: SessionHistoryPagination = {
         hasOlder: window.hasOlder,
@@ -575,7 +669,16 @@ export async function buildSnapshotFromSessionEvents(
         historyPagination,
     });
 
-    return attachRunningTurn(withHistory, runningTurn);
+    const tip = await resolveSessionTip({
+        server,
+        sessionId,
+        itemsAsc: window.itemsAsc,
+    });
+    // The tip has no turn.done in this window, so ingestion never folded its
+    // input. Apply it here so answered approvals / ask-user prompts are not
+    // restored as pending after a refresh.
+    applyUserToolResponsesToFold(withHistory.fold, tip.continuationInput);
+    return attachRunningTurn(withHistory, tip.runningTurn);
 }
 
 /**
@@ -1276,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,