@truefoundry/assistant-ui-runtime 0.1.1 → 0.1.3-rc.1

package/src/private/useDraftAgentSpec.ts

@@ -19,13 +19,22 @@ export type UseDraftAgentSpecOptions = {
     onError?: ((error: unknown) => void) | undefined;
 };
 
+export type UseDraftAgentSpecResult = {
+    agentSpec: AgentSpec | null;
+    draftSessionId: string | undefined;
+    isSpecSyncing: boolean;
+    specError: unknown | null;
+    updateAgentSpec: (update: AgentSpecUpdate) => void;
+    takeTurnHeaderTimestamp: () => Promise<string | undefined>;
+};
+
 export function useDraftAgentSpec({
     draftSessionId,
     draftBridge,
     defaultAgentSpec,
     onAgentSpecChange,
     onError,
-}: UseDraftAgentSpecOptions) {
+}: UseDraftAgentSpecOptions): UseDraftAgentSpecResult {
     const enabled = draftBridge != null;
     const [agentSpec, setAgentSpec] = useState<AgentSpec>(defaultAgentSpec);
     const [isSpecSyncing, setIsSpecSyncing] = useState(false);
@@ -38,6 +47,35 @@ export function useDraftAgentSpec({
     const syncGenerationRef = useRef(0);
     const loadedDraftIdRef = useRef<string | undefined>(undefined);
     const localDirtyRef = useRef(false);
+    const lastUpdatedAtRef = useRef<string | undefined>(undefined);
+    const pendingFlushRef = useRef<(() => Promise<void>) | undefined>(undefined);
+    const inFlightFlushRef = useRef<Promise<void> | undefined>(undefined);
+    const activeDraftIdRef = useRef(draftSessionId);
+
+    // Invalidate sync state from the previous draft so a stale pending flush or
+    // stored updatedAt can't update the wrong draft or leak into the next
+    // turn's header. Bumping the generation makes any in-flight sync a no-op.
+    useEffect(() => {
+        const previousDraftId = activeDraftIdRef.current;
+        if (previousDraftId === draftSessionId) {
+            return;
+        }
+        activeDraftIdRef.current = draftSessionId;
+        syncGenerationRef.current++;
+        if (syncTimeoutRef.current != null) {
+            clearTimeout(syncTimeoutRef.current);
+            syncTimeoutRef.current = undefined;
+        }
+        pendingFlushRef.current = undefined;
+        inFlightFlushRef.current = undefined;
+        lastUpdatedAtRef.current = undefined;
+        setIsSpecSyncing(false);
+        // Dirty edits made against a previous draft must not be replayed onto
+        // the new one. Keep them only for lazy creation (undefined -> id).
+        if (previousDraftId != null) {
+            localDirtyRef.current = false;
+        }
+    }, [draftSessionId]);
 
     useEffect(() => {
         if (!enabled || draftBridge == null) {
@@ -93,10 +131,12 @@ export function useDraftAgentSpec({
             }
             setIsSpecSyncing(true);
             try {
-                await draftBridge.syncAgentSpec(draftId, spec);
+                const updatedAt = await draftBridge.syncAgentSpec(draftId, spec);
                 if (generation !== syncGenerationRef.current) {
                     return;
                 }
+                lastUpdatedAtRef.current =
+                    updatedAt || new Date().toISOString();
                 setSpecError(null);
                 onAgentSpecChange?.(spec);
             } catch (error) {
@@ -119,8 +159,20 @@ export function useDraftAgentSpec({
                 clearTimeout(syncTimeoutRef.current);
             }
             const generation = ++syncGenerationRef.current;
+            const flush = () => {
+                pendingFlushRef.current = undefined;
+                syncTimeoutRef.current = undefined;
+                const promise = flushSpecSync(draftId, spec, generation).finally(() => {
+                    if (inFlightFlushRef.current === promise) {
+                        inFlightFlushRef.current = undefined;
+                    }
+                });
+                inFlightFlushRef.current = promise;
+                return promise;
+            };
+            pendingFlushRef.current = flush;
             syncTimeoutRef.current = setTimeout(() => {
-                void flushSpecSync(draftId, spec, generation);
+                void flush();
             }, SPEC_SYNC_DEBOUNCE_MS);
         },
         [flushSpecSync],
@@ -129,6 +181,29 @@ export function useDraftAgentSpec({
     const scheduleSpecSyncRef = useRef(scheduleSpecSync);
     scheduleSpecSyncRef.current = scheduleSpecSync;
 
+    const flushPendingSpecSyncNow = useCallback(async () => {
+        if (syncTimeoutRef.current != null) {
+            clearTimeout(syncTimeoutRef.current);
+            syncTimeoutRef.current = undefined;
+        }
+        const pending = pendingFlushRef.current;
+        if (pending != null) {
+            pendingFlushRef.current = undefined;
+            await pending();
+            return;
+        }
+        if (inFlightFlushRef.current != null) {
+            await inFlightFlushRef.current;
+        }
+    }, []);
+
+    const takeTurnHeaderTimestamp = useCallback(async () => {
+        await flushPendingSpecSyncNow();
+        const updatedAt = lastUpdatedAtRef.current;
+        lastUpdatedAtRef.current = undefined;
+        return updatedAt;
+    }, [flushPendingSpecSyncNow]);
+
     useEffect(
         () => () => {
             if (syncTimeoutRef.current != null) {
@@ -160,6 +235,7 @@ export function useDraftAgentSpec({
             isSpecSyncing: enabled ? isSpecSyncing : false,
             specError: enabled ? specError : null,
             updateAgentSpec,
+            takeTurnHeaderTimestamp,
         }),
         [
             agentSpec,
@@ -167,6 +243,7 @@ export function useDraftAgentSpec({
             enabled,
             isSpecSyncing,
             specError,
+            takeTurnHeaderTimestamp,
             updateAgentSpec,
         ],
     );

package/src/sessionSnapshot.ts

@@ -3,6 +3,7 @@ import type {
     Turn,
     TurnCreatedEvent,
     TurnDoneEvent,
+    TurnEvent,
     TurnInputItem,
 } from "truefoundry-gateway-sdk/agents";
 
@@ -12,6 +13,19 @@ import type { StoredApprovalDecision } from "./toolApproval.js";
 import type { StoredToolResponse } from "./toolResponse.js";
 import type { TurnStreamUpdate } from "./turnStreamUpdate.js";
 
+/** Session-level event item from `AgentSession.listEvents`. */
+export type GatewaySessionEventItem = {
+    turnId: string;
+    event: TurnCreatedEvent | TurnDoneEvent | TurnEvent;
+};
+
+/** Cursor for fetching older `listEvents` pages (scroll-up history). */
+export type SessionHistoryPagination = {
+    /** Token for the next older page; omit when exhausted. */
+    olderPageToken?: string;
+    hasOlder: boolean;
+};
+
 /** Turn metadata retained for cross-turn projection and subsequent required-action replay. */
 export type SessionTurnRecord = Pick<
     Turn,
@@ -55,6 +69,12 @@ export type SessionSnapshot = {
     requiredActions: RequiredActionsOverlay;
     runningTurn?: Turn;
     unstable_resume?: boolean;
+    /**
+     * Chronological `listEvents` items loaded so far (for prepend-on-scroll rebuild).
+     * Live stream commits are not appended here — they live in `turns` / `fold`.
+     */
+    historyEvents?: readonly GatewaySessionEventItem[];
+    historyPagination?: SessionHistoryPagination;
 };
 
 export type ProjectSessionMessagesOptions = {
@@ -110,7 +130,7 @@ export function sessionEventsToSessionRecord(
 /** Returns a new snapshot wrapper; fold maps may be mutated in place before calling. */
 export function replaceSessionSnapshot(
     snapshot: SessionSnapshot,
-    patch: Partial<Omit<SessionSnapshot, "fold" | "requiredActions">> & {
+    patch: Partial<Omit<SessionSnapshot, "requiredActions">> & {
         requiredActions?: RequiredActionsOverlay;
     },
 ): SessionSnapshot {

package/src/sessions.ts

@@ -2,15 +2,15 @@ import type {
     AgentSession,
     AgentSessionClient,
 } from "truefoundry-gateway-sdk/agents";
-import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
+import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
 
 import { bindDraftAgentSession } from "./private/bindDraftAgentSession.js";
 
 const inflightBySessionId = new Map<string, Promise<AgentSession>>();
 
 export type GetSessionOptions = {
-    /** When set, validates the draft and binds turns to `/agents/sessions/{id}/turns`. */
-    draftGateway?: TrueFoundryGateway;
+    /** When set, validates the draft and binds turns via PrivateAgentSessionClient. */
+    privateClient?: PrivateAgentSessionClient;
 };
 
 /** `sessionId` is the assistant-ui thread `remoteId` from `RemoteThreadListAdapter.initialize`. */
@@ -19,8 +19,8 @@ export function getSession(
     sessionId: string,
     options?: GetSessionOptions,
 ): Promise<AgentSession> {
-    if (options?.draftGateway != null) {
-        return bindDraftAgentSession(client, options.draftGateway, sessionId);
+    if (options?.privateClient != null) {
+        return bindDraftAgentSession(options.privateClient, sessionId);
     }
 
     let inflight = inflightBySessionId.get(sessionId);

package/src/streamTurn.test.ts

@@ -193,6 +193,40 @@ describe("streamTurn", () => {
             expect(execute).not.toHaveBeenCalled();
             expect(updates).toEqual([]);
         });
+
+        it("forwards headers to turn.execute request options", async () => {
+            const execute = vi.fn(() => (async function* () {})());
+            const prepareTurn = vi.fn(() => ({ execute }));
+            const session = {
+                prepareTurn,
+                cancel: vi.fn().mockResolvedValue(undefined),
+            } as unknown as AgentSession;
+            const abortSignal = new AbortController().signal;
+
+            await collectUpdates(
+                streamTurnContent(
+                    session,
+                    new PeerThreadFoldState(),
+                    {
+                        userMessage: "hello",
+                        headers: {
+                            "x-tfy-session-last-updated-at": "2026-06-30T10:00:00.000Z",
+                        },
+                    },
+                    abortSignal,
+                ),
+            );
+
+            expect(execute).toHaveBeenCalledWith(
+                { stream: true },
+                {
+                    abortSignal,
+                    headers: {
+                        "x-tfy-session-last-updated-at": "2026-06-30T10:00:00.000Z",
+                    },
+                },
+            );
+        });
     });
 
     describe("resumeTurnStream", () => {

package/src/streamTurn.ts

@@ -21,6 +21,8 @@ export type StreamTurnOptions = {
      * root turn (no `previousTurnId` field).
      */
     previousTurnId?: string | null;
+    /** Extra headers for the createTurn request (`execute` request options). */
+    headers?: Record<string, string>;
 };
 
 function buildTurnInput(options: StreamTurnOptions): TurnInputItem[] {
@@ -68,7 +70,13 @@ export async function* streamTurnContent(
 
     try {
         yield* streamTurnEvents(
-            turn.execute({ stream: true }, { abortSignal }),
+            turn.execute(
+                { stream: true },
+                {
+                    abortSignal,
+                    ...(options.headers != null ? { headers: options.headers } : {}),
+                },
+            ),
             foldState,
             groupRootBaseline,
         );

package/src/truefoundryExtras.ts

@@ -28,6 +28,9 @@ export type TrueFoundryRuntimeExtras = {
     cancel: () => Promise<void>;
     resetFromTurn: (turnId: string) => Promise<void>;
     reload: () => void;
+    hasOlderHistory: boolean;
+    isLoadingOlderHistory: boolean;
+    loadOlderHistory: () => Promise<void>;
     draft: TrueFoundryDraftRuntimeExtras | null;
 };
 

package/src/truefoundryOwnedSessionsThreadListAdapter.test.ts

@@ -0,0 +1,100 @@
+import { describe, expect, it, vi } from "vitest";
+
+import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
+
+import { createTrueFoundryOwnedSessionsThreadListAdapter } from "./truefoundryOwnedSessionsThreadListAdapter.js";
+
+function mockNamedSession(id: string, title: string, updatedAt: string) {
+    return {
+        type: "session" as const,
+        id,
+        agentName: "my-agent",
+        title,
+        createdBySubject: { type: "user" as const, id: "u1" },
+        createdAt: updatedAt,
+        updatedAt,
+    };
+}
+
+function mockDraftSession(id: string, title: string | undefined, updatedAt: string) {
+    return {
+        type: "session/draft" as const,
+        id,
+        agentSpec: { model: { name: "anthropic/claude-sonnet-4-6" } },
+        title,
+        createdBySubject: { type: "user" as const, id: "u1" },
+        createdAt: updatedAt,
+        updatedAt,
+    };
+}
+
+describe("createTrueFoundryOwnedSessionsThreadListAdapter", () => {
+    it("lists owned sessions (named + draft) with pagination cursor", async () => {
+        const listOwnedSessions = vi.fn().mockResolvedValue({
+            data: [
+                mockNamedSession("s1", "Named chat", "2026-06-30T12:00:00.000Z"),
+                mockDraftSession("d1", "Draft chat", "2026-06-30T11:00:00.000Z"),
+            ],
+            response: {
+                pagination: { nextPageToken: "page-2", limit: 20 },
+            },
+        });
+        const privateClient = {
+            listOwnedSessions,
+            getDraftSession: vi.fn(),
+        } as unknown as PrivateAgentSessionClient;
+
+        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ privateClient });
+        const result = await adapter.list();
+
+        expect(listOwnedSessions).toHaveBeenCalledWith(
+            expect.objectContaining({
+                limit: 20,
+                pageToken: undefined,
+                startTimestamp: expect.any(String),
+            }),
+        );
+        expect(result.threads).toEqual([
+            {
+                status: "regular",
+                remoteId: "s1",
+                title: "Named chat",
+                lastMessageAt: new Date("2026-06-30T12:00:00.000Z"),
+            },
+            {
+                status: "regular",
+                remoteId: "d1",
+                title: "Draft chat",
+                lastMessageAt: new Date("2026-06-30T11:00:00.000Z"),
+            },
+        ]);
+        expect(result.nextCursor).toBe("page-2");
+    });
+
+    it("falls back to model name for untitled drafts", async () => {
+        const listOwnedSessions = vi.fn().mockResolvedValue({
+            data: [mockDraftSession("d1", undefined, "2026-06-30T11:00:00.000Z")],
+            response: { pagination: { limit: 20 } },
+        });
+        const privateClient = {
+            listOwnedSessions,
+            getDraftSession: vi.fn(),
+        } as unknown as PrivateAgentSessionClient;
+
+        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ privateClient });
+        const result = await adapter.list();
+
+        expect(result.threads[0]?.title).toBe("anthropic/claude-sonnet-4-6");
+    });
+
+    it("throws on initialize because the adapter is read-only", async () => {
+        const privateClient = {
+            listOwnedSessions: vi.fn(),
+            getDraftSession: vi.fn(),
+        } as unknown as PrivateAgentSessionClient;
+
+        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ privateClient });
+
+        await expect(adapter.initialize("local")).rejects.toThrow(/read-only/);
+    });
+});

package/src/truefoundryOwnedSessionsThreadListAdapter.ts

@@ -0,0 +1,77 @@
+import type { RemoteThreadListAdapter } from "@assistant-ui/core";
+import type { AgentSession } from "truefoundry-gateway-sdk/agents";
+import type {
+    AgentDraftSession,
+    PrivateAgentSessionClient,
+} from "truefoundry-gateway-sdk/agents/private";
+
+import { draftSessionTitle } from "./private/agentSpec.js";
+import { sessionListStartTimestamp } from "./sessionListStartTimestamp.js";
+
+const THREAD_LIST_PAGE_SIZE = 20;
+
+function ownedSessionTitle(session: AgentSession | AgentDraftSession): string {
+    if (session.type === "session/draft") {
+        return draftSessionTitle(session);
+    }
+    return session.title ?? session.agentName;
+}
+
+/**
+ * Read-only thread-list adapter backed by `PrivateAgentSessionClient.listOwnedSessions`.
+ * Returns every session the caller owns (named + draft), newest first.
+ */
+export function createTrueFoundryOwnedSessionsThreadListAdapter(options: {
+    privateClient: PrivateAgentSessionClient;
+}): RemoteThreadListAdapter {
+    const { privateClient } = options;
+
+    return {
+        async list({ after } = {}) {
+            const page = await privateClient.listOwnedSessions({
+                limit: THREAD_LIST_PAGE_SIZE,
+                pageToken: after,
+                startTimestamp: sessionListStartTimestamp(),
+            });
+            const threads = page.data.map((session) => ({
+                status: "regular" as const,
+                remoteId: session.id,
+                title: ownedSessionTitle(session),
+                lastMessageAt: new Date(session.updatedAt),
+            }));
+            return {
+                threads,
+                nextCursor: page.response.pagination.nextPageToken ?? undefined,
+            };
+        },
+
+        async initialize() {
+            throw new Error(
+                "Owned sessions history adapter is read-only; create sessions via a named or draft runtime.",
+            );
+        },
+
+        async fetch(remoteId) {
+            // PrivateAgentSessionClient only exposes get for drafts. Named sessions
+            // surface through listOwnedSessions; fetch falls back to draft get.
+            const draft = await privateClient.getDraftSession({
+                draftSessionId: remoteId,
+            });
+            return {
+                status: "regular" as const,
+                remoteId: draft.id,
+                title: draftSessionTitle(draft),
+                lastMessageAt: new Date(draft.updatedAt),
+            };
+        },
+
+        async rename() {},
+        async archive() {},
+        async unarchive() {},
+        async delete() {},
+
+        async generateTitle() {
+            return new ReadableStream();
+        },
+    };
+}

package/src/types.ts

@@ -7,7 +7,7 @@ import type {
     SpeechSynthesisAdapter,
 } from "@assistant-ui/core";
 import type { AgentSessionClient } from "truefoundry-gateway-sdk/agents";
-import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
+import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
 
 import type { AgentSpec } from "./private/agentSpec.js";
 
@@ -47,13 +47,16 @@ export type UseTrueFoundryAgentRuntimeOptions = TrueFoundryAgentRuntimeBaseOptio
     agent?: TrueFoundryAgentConfig | undefined;
     /** Legacy named-agent shorthand. Prefer `agent: { mode: "named", agentName }`. */
     agentName?: string | undefined;
-    /** Required when `agent.mode === "draft"`. */
-    gateway?: TrueFoundryGateway | undefined;
+    /**
+     * Required when `agent.mode === "draft"`. Also used for sandbox file downloads
+     * (named or draft).
+     */
+    privateClient?: PrivateAgentSessionClient | undefined;
 };
 
 export type ResolvedTrueFoundryAgentRuntimeOptions = TrueFoundryAgentRuntimeBaseOptions & {
     agent: TrueFoundryAgentConfig;
-    gateway?: TrueFoundryGateway | undefined;
+    privateClient?: PrivateAgentSessionClient | undefined;
 };
 
 export function resolveTrueFoundryAgentConfig(
@@ -78,15 +81,15 @@ export function resolveTrueFoundryAgentRuntimeOptions(
 ): ResolvedTrueFoundryAgentRuntimeOptions {
     const agent = resolveTrueFoundryAgentConfig(options);
 
-    if (agent.mode === "draft" && options.gateway == null) {
+    if (agent.mode === "draft" && options.privateClient == null) {
         throw new Error(
-            "Draft agent mode requires a `gateway` TrueFoundryGateway client.",
+            "Draft agent mode requires a `privateClient` PrivateAgentSessionClient.",
         );
     }
 
     return {
         ...options,
         agent,
-        gateway: options.gateway,
+        privateClient: options.privateClient,
     };
 }

package/src/useTrueFoundryAgentMessages.test.tsx

@@ -3,7 +3,7 @@ import type { ThreadMessage } from "@assistant-ui/core";
 import { act, renderHook, waitFor } from "@testing-library/react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { AgentSessionClient, Turn } from "truefoundry-gateway-sdk/agents";
-import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
+import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
 
 import { ROOT_THREAD_ID } from "./constants.js";
 import { collectPendingToolResponses } from "./collectPending.js";
@@ -329,6 +329,56 @@ describe("useTrueFoundryAgentMessages", () => {
         expect(loadSessionSnapshot).not.toHaveBeenCalled();
     });
 
+    it("sendTurn forwards getTurnHeaders only when they resolve to a value", async () => {
+        const getTurnHeaders = vi
+            .fn()
+            .mockResolvedValueOnce({
+                "x-tfy-session-last-updated-at": "2026-06-30T12:00:00.000Z",
+            })
+            .mockResolvedValueOnce(undefined);
+
+        const { result } = renderHook(() =>
+            useTrueFoundryAgentMessages({
+                client: mockClient,
+                sessionId: "session-1",
+                getTurnHeaders,
+            }),
+        );
+
+        await waitFor(() => expect(result.current.isLoading).toBe(false));
+
+        await act(async () => {
+            await result.current.sendTurn({ userMessage: "first" });
+        });
+
+        expect(getTurnHeaders).toHaveBeenCalledTimes(1);
+        expect(streamTurnContent).toHaveBeenCalledWith(
+            expect.anything(),
+            expect.any(PeerThreadFoldState),
+            {
+                userMessage: "first",
+                headers: {
+                    "x-tfy-session-last-updated-at": "2026-06-30T12:00:00.000Z",
+                },
+            },
+            expect.any(AbortSignal),
+            expect.any(Array),
+        );
+
+        await act(async () => {
+            await result.current.sendTurn({ userMessage: "second" });
+        });
+
+        expect(getTurnHeaders).toHaveBeenCalledTimes(2);
+        expect(streamTurnContent).toHaveBeenLastCalledWith(
+            expect.anything(),
+            expect.any(PeerThreadFoldState),
+            { userMessage: "second" },
+            expect.any(AbortSignal),
+            expect.any(Array),
+        );
+    });
+
     it("loads converted session history on mount", async () => {
         vi.mocked(loadSessionSnapshot).mockResolvedValue(
             snapshotWithUserTurn("hello"),
@@ -921,16 +971,16 @@ describe("useTrueFoundryAgentMessages", () => {
         expect(loadSessionSnapshot).toHaveBeenCalledTimes(1);
     });
 
-    it("cancel resolves the session through the draft gateway in draft mode", async () => {
+    it("cancel resolves the session through the private client in draft mode", async () => {
         const cancel = vi.fn().mockResolvedValue(undefined);
         vi.mocked(getSession).mockResolvedValue({ cancel } as never);
-        const draftGateway = {} as TrueFoundryGateway;
+        const privateClient = {} as PrivateAgentSessionClient;
 
         const { result } = renderHook(() =>
             useTrueFoundryAgentMessages({
                 client: mockClient,
                 sessionId: "draft-session-1",
-                draftGateway,
+                privateClient,
             }),
         );
         await waitFor(() => expect(result.current.isLoading).toBe(false));
@@ -941,7 +991,7 @@ describe("useTrueFoundryAgentMessages", () => {
 
         expect(cancel).toHaveBeenCalled();
         expect(getSession).toHaveBeenCalledWith(mockClient, "draft-session-1", {
-            draftGateway,
+            privateClient,
         });
     });
 });