@truefoundry/assistant-ui-runtime 0.1.2 → 0.1.3-rc.2

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";
@@ -971,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));
@@ -991,7 +991,7 @@ describe("useTrueFoundryAgentMessages", () => {
 
         expect(cancel).toHaveBeenCalled();
         expect(getSession).toHaveBeenCalledWith(mockClient, "draft-session-1", {
-            draftGateway,
+            privateClient,
         });
     });
 });

package/src/useTrueFoundryAgentMessages.ts

@@ -11,7 +11,7 @@ import type {
     TurnInputItem,
     TurnStateDone,
 } 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 {
@@ -68,8 +68,8 @@ export type UseTrueFoundryAgentMessagesOptions = {
     }>;
     /** Maps a thread `remoteId` to the gateway session id used for turns. */
     resolveConversationSessionId?: (remoteId: string) => Promise<string>;
-    /** When set, turns bind to `/agents/sessions/{draftSessionId}/turns` after draft validation. */
-    draftGateway?: TrueFoundryGateway;
+    /** When set, turns bind via PrivateAgentSessionClient.getDraftSession. */
+    privateClient?: PrivateAgentSessionClient;
     /**
      * Optional per-turn headers for createTurn. Invoked once per `sendTurn` after
      * the session is resolved; return value is forwarded to `turn.execute`.
@@ -294,13 +294,13 @@ export function useTrueFoundryAgentMessages({
     onError,
     initializeSession,
     resolveConversationSessionId,
-    draftGateway,
+    privateClient,
     getTurnHeaders,
 }: UseTrueFoundryAgentMessagesOptions) {
     const sessionOptions = useMemo<GetSessionOptions | undefined>(
         () =>
-            draftGateway != null ? { draftGateway } : undefined,
-        [draftGateway],
+            privateClient != null ? { privateClient } : undefined,
+        [privateClient],
     );
 
     const [snapshot, setSnapshot] = useState<SessionSnapshot>(createEmptySessionSnapshot);
@@ -353,7 +353,14 @@ export function useTrueFoundryAgentMessages({
     const runStream = useCallback(
         (
             createStream: (signal: AbortSignal) => AsyncGenerator<TurnStreamUpdate>,
-            turnId: string,
+            /**
+             * A mutable ref whose `.current` is the turn ID to use for
+             * `activeStream.turnId`. Callers that capture the gateway turn ID
+             * via `onTurnIdAvailable` update this ref in-place so that both the
+             * pending-update flush and `commitActiveStream` always see the real
+             * gateway ID rather than the locally-generated optimistic one.
+             */
+            turnIdRef: { current: string },
             isContinuation: boolean,
         ): Promise<void> => {
             const streamGeneration = ++streamGenerationRef.current;
@@ -368,7 +375,6 @@ export function useTrueFoundryAgentMessages({
                 // message tree (UI hang). The buffer belongs to this stream only.
                 let pendingStreamUpdate: {
                     update: TurnStreamUpdate;
-                    turnId: string;
                     isContinuation: boolean;
                 } | null = null;
                 let streamUpdateRaf: number | null = null;
@@ -383,12 +389,14 @@ export function useTrueFoundryAgentMessages({
                     ) {
                         return;
                     }
-                    const { update, turnId: pendingTurnId, isContinuation: pendingIsContinuation } =
-                        pending;
+                    const { update, isContinuation: pendingIsContinuation } = pending;
                     setSnapshot((prev) =>
                         replaceSessionSnapshot(prev, {
                             activeStream: {
-                                turnId: pendingTurnId,
+                                // Read from the ref so we always use the latest ID,
+                                // including any gateway ID that arrived after the RAf
+                                // was scheduled.
+                                turnId: turnIdRef.current,
                                 update,
                                 isContinuation: pendingIsContinuation,
                             },
@@ -397,7 +405,7 @@ export function useTrueFoundryAgentMessages({
                 };
 
                 const applyStreamUpdate = (update: TurnStreamUpdate) => {
-                    pendingStreamUpdate = { update, turnId, isContinuation };
+                    pendingStreamUpdate = { update, isContinuation };
                     if (streamUpdateRaf == null) {
                         streamUpdateRaf = requestAnimationFrame(flushPendingStreamUpdate);
                     }
@@ -540,7 +548,7 @@ export function useTrueFoundryAgentMessages({
                             undefined,
                             loadedSnapshot.groupRootBaseline,
                         ),
-                    turn.id,
+                    { current: turn.id },
                     isContinuation,
                 ).catch(() => undefined);
             }
@@ -589,6 +597,11 @@ export function useTrueFoundryAgentMessages({
                     ? continuationTurnId
                     : generateId();
 
+            // Mutable ref so runStream always reads the latest ID. For new
+            // user-message turns the local `generateId()` value is replaced
+            // with the gateway-assigned ID once the first SSE event arrives.
+            const turnIdRef = { current: turnId };
+
             if ("inputs" in options) {
                 applyUserToolResponsesToFold(
                     snapshotRef.current.fold,
@@ -681,9 +694,26 @@ export function useTrueFoundryAgentMessages({
                         },
                         signal,
                         groupRootBaseline,
+                        // Rename the optimistic local ID to the gateway turn ID
+                        // so that edit/retry can resolve the turn via the gateway.
+                        (gatewayTurnId) => {
+                            const oldId = turnIdRef.current;
+                            if (gatewayTurnId === oldId) return;
+                            turnIdRef.current = gatewayTurnId;
+                            // Rename in the ref immediately so any synchronous read
+                            // (e.g. commitActiveStream) sees the correct ID.
+                            const renamePendingUser = (prev: SessionSnapshot): SessionSnapshot => {
+                                if (prev.pendingUser?.turnId !== oldId) return prev;
+                                return replaceSessionSnapshot(prev, {
+                                    pendingUser: { ...prev.pendingUser, turnId: gatewayTurnId },
+                                });
+                            };
+                            snapshotRef.current = renamePendingUser(snapshotRef.current);
+                            setSnapshot(renamePendingUser);
+                        },
                     );
                 },
-                turnId,
+                turnIdRef,
                 isContinuation,
             );
         },
@@ -784,7 +814,7 @@ export function useTrueFoundryAgentMessages({
                     undefined,
                     snapshotRef.current.groupRootBaseline,
                 ),
-            turn.id,
+            { current: turn.id },
             true,
         );
     }, [runStream]);

package/src/useTrueFoundryAgentRuntime.ts

@@ -30,6 +30,7 @@ import {
     createDraftSessionBridge,
     DRAFT_SESSION_LAST_UPDATED_AT_HEADER,
 } from "./private/draftSessionBridge.js";
+import { getGatewayFromPrivateClient } from "./private/getGatewayFromPrivateClient.js";
 import { MCP_AUTH_RESUME_RUN_CUSTOM_KEY } from "./mcpAuth.js";
 import { createTrueFoundryDraftThreadListAdapter } from "./private/truefoundryDraftThreadListAdapter.js";
 import { trueFoundryExtras } from "./truefoundryExtras.js";
@@ -46,7 +47,7 @@ function useTrueFoundryAgentRuntimeImpl(
     const {
         client,
         agent,
-        gateway,
+        privateClient,
         adapters,
         onError,
         listEventsConcurrency,
@@ -54,8 +55,8 @@ function useTrueFoundryAgentRuntimeImpl(
     } = options;
 
     const draftBridgeRef = useRef(
-        agent.mode === "draft" && gateway != null
-            ? createDraftSessionBridge(gateway)
+        agent.mode === "draft" && privateClient != null
+            ? createDraftSessionBridge(privateClient)
             : null,
     );
 
@@ -125,7 +126,7 @@ function useTrueFoundryAgentRuntimeImpl(
         listEventsConcurrency,
         onError,
         initializeSession,
-        draftGateway: agent.mode === "draft" ? gateway : undefined,
+        privateClient: agent.mode === "draft" ? privateClient : undefined,
         getTurnHeaders: agent.mode === "draft" ? getTurnHeaders : undefined,
     });
 
@@ -151,9 +152,9 @@ function useTrueFoundryAgentRuntimeImpl(
 
     const downloadSandboxFile = useCallback(
         async (path: string) => {
-            if (gateway == null) {
+            if (privateClient == null) {
                 const error = new Error(
-                    "Downloading a sandbox file requires a `gateway` TrueFoundryGateway client.",
+                    "Downloading a sandbox file requires a `privateClient` PrivateAgentSessionClient.",
                 );
                 onError?.(error);
                 throw error;
@@ -163,10 +164,12 @@ function useTrueFoundryAgentRuntimeImpl(
                 onError?.(error);
                 throw error;
             }
+            // Same gateway.agents.downloadSandboxFile path as before, via the wrapped client.
+            const gateway = getGatewayFromPrivateClient(privateClient);
             const response = await gateway.agents.downloadSandboxFile(sandboxId, { path });
             return await response.blob();
         },
-        [gateway, sandboxId, onError],
+        [privateClient, sandboxId, onError],
     );
 
     const draftExtras = useMemo(() => {
@@ -264,7 +267,7 @@ export function useTrueFoundryAgentRuntime(options: UseTrueFoundryAgentRuntimeOp
     // causing `threadListAdapter` to be a new object each render.  We compute `resolved`
     // unconditionally and stabilize individual downstream values with primitives/refs.
     const resolved = resolveTrueFoundryAgentRuntimeOptions(options);
-    const { client, agent, gateway } = resolved;
+    const { client, agent, privateClient } = resolved;
 
     const pendingAgentSpecRef = useRef<AgentSpec | undefined>(
         agent.mode === "draft" ? agent.defaultAgentSpec : undefined,
@@ -276,14 +279,14 @@ export function useTrueFoundryAgentRuntime(options: UseTrueFoundryAgentRuntimeOp
     const namedAgentName = agent.mode === "named" ? agent.agentName : undefined;
     const threadListAdapter = useMemo(() => {
         if (agentMode === "draft") {
-            if (gateway == null) {
+            if (privateClient == null) {
                 throw new Error(
-                    "Draft agent mode requires a `gateway` TrueFoundryGateway client.",
+                    "Draft agent mode requires a `privateClient` PrivateAgentSessionClient.",
                 );
             }
             const draftAgent = agent as Extract<typeof agent, { mode: "draft" }>;
             return createTrueFoundryDraftThreadListAdapter({
-                gateway,
+                privateClient,
                 defaultAgentSpec: draftAgent.defaultAgentSpec,
                 getAgentSpec: () => pendingAgentSpecRef.current ?? draftAgent.defaultAgentSpec,
             });
@@ -293,7 +296,7 @@ export function useTrueFoundryAgentRuntime(options: UseTrueFoundryAgentRuntimeOp
             agentName: namedAgentName!,
         });
         // eslint-disable-next-line react-hooks/exhaustive-deps
-    }, [agentMode, namedAgentName, client, gateway]);
+    }, [agentMode, namedAgentName, client, privateClient]);
 
     return useRemoteThreadListRuntime({
         allowNesting: true,

package/src/agentSessionModule.d.ts

@@ -1,37 +0,0 @@
-declare module "truefoundry-gateway-sdk/dist/esm/agents/AgentSession.mjs" {
-    import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
-    import type * as TrueFoundryGatewayApi from "truefoundry-gateway-sdk";
-    import type { PreparedTurn } from "truefoundry-gateway-sdk/agents";
-    import type * as core from "truefoundry-gateway-sdk/core";
-
-    export class AgentSession implements TrueFoundryGatewayApi.Session {
-        readonly id: string;
-        readonly agentName: string;
-        readonly title?: string;
-        readonly createdBySubject: TrueFoundryGatewayApi.Subject;
-        readonly createdAt: string;
-        readonly updatedAt: string;
-        constructor(session: TrueFoundryGatewayApi.Session, client: TrueFoundryGateway);
-        prepareTurn(opts?: {
-            input?: TrueFoundryGatewayApi.TurnInputItem[];
-            previousTurnId?: TrueFoundryGatewayApi.PreviousTurnIdInput;
-        }): PreparedTurn;
-        listTurns(
-            opts?: TrueFoundryGatewayApi.agents.SessionsListTurnsRequest,
-            requestOptions?: unknown,
-        ): Promise<core.Page<unknown, TrueFoundryGatewayApi.ListTurnsResponse>>;
-        /**
-         * Paginated session events across turns (newest first); subscribe to a
-         * running turn for live events.
-         *
-         * @param opts.lastTurnId - Newest turn in the listing window (initial load only).
-         *   Omit to use the session last turn. Running-turn events are excluded.
-         */
-        listEvents(
-            opts?: TrueFoundryGatewayApi.agents.SessionsListEventsRequest,
-            requestOptions?: unknown,
-        ): Promise<core.Page<TrueFoundryGatewayApi.SessionEventItem, TrueFoundryGatewayApi.ListSessionEventsResponse>>;
-        getTurn(opts: { turnId: string }, requestOptions?: unknown): Promise<unknown>;
-        cancel(requestOptions?: unknown): Promise<void>;
-    }
-}