@truefoundry/assistant-ui-runtime 0.1.4 → 0.1.6-rc.0

package/src/useTrueFoundryAgentMessages.test.tsx

@@ -2,8 +2,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 { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
+import type { AgentChatServer, Turn } from "./server/index.js";
 
 import { ROOT_THREAD_ID } from "./constants.js";
 import { collectPendingToolResponses } from "./collectPending.js";
@@ -18,7 +17,6 @@ import {
     replaceSessionSnapshot,
     type SessionSnapshot,
 } from "./sessionSnapshot.js";
-import { getSession } from "./sessions.js";
 import { resumeTurnStream, streamTurnContent } from "./streamTurn.js";
 import {
     messageHasPendingApprovals,
@@ -32,10 +30,6 @@ import {
 } from "./toolResponse.js";
 import { useTrueFoundryAgentMessages } from "./useTrueFoundryAgentMessages.js";
 
-vi.mock("./sessions.js", () => ({
-    getSession: vi.fn(),
-}));
-
 vi.mock("./loadSessionSnapshot.js", () => ({
     loadSessionSnapshot: vi.fn(),
 }));
@@ -50,7 +44,9 @@ vi.mock("./convertTurnMessages.js", async (importOriginal) => {
     return actual;
 });
 
-const mockClient = {} as AgentSessionClient;
+const mockServer = {
+    cancelSession: vi.fn().mockResolvedValue(undefined),
+} as unknown as AgentChatServer;
 
 function snapshotWithAssistantMessage(
     message: Extract<ThreadMessage, { role: "assistant" }>,
@@ -282,7 +278,7 @@ async function* singleUpdateStream() {
 describe("useTrueFoundryAgentMessages", () => {
     beforeEach(() => {
         vi.clearAllMocks();
-        vi.mocked(getSession).mockResolvedValue({} as never);
+        vi.mocked(mockServer.cancelSession).mockResolvedValue(undefined);
         vi.mocked(loadSessionSnapshot).mockResolvedValue(createEmptySessionSnapshot());
         vi.mocked(streamTurnContent).mockReturnValue(singleUpdateStream());
         vi.mocked(resumeTurnStream).mockReturnValue(singleUpdateStream());
@@ -294,12 +290,11 @@ describe("useTrueFoundryAgentMessages", () => {
 
     it("clears messages when sessionId is undefined", async () => {
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: undefined }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: undefined }),
         );
 
         await waitFor(() => expect(result.current.isLoading).toBe(false));
         expect(result.current.messages).toEqual([]);
-        expect(getSession).not.toHaveBeenCalled();
         expect(loadSessionSnapshot).not.toHaveBeenCalled();
     });
 
@@ -311,7 +306,7 @@ describe("useTrueFoundryAgentMessages", () => {
 
         const { result } = renderHook(() =>
             useTrueFoundryAgentMessages({
-                client: mockClient,
+                server: mockServer,
                 sessionId: undefined,
                 initializeSession,
             }),
@@ -324,7 +319,6 @@ describe("useTrueFoundryAgentMessages", () => {
         });
 
         expect(initializeSession).toHaveBeenCalledOnce();
-        expect(getSession).toHaveBeenCalledWith(mockClient, "session-new");
         expect(streamTurnContent).toHaveBeenCalled();
         expect(loadSessionSnapshot).not.toHaveBeenCalled();
     });
@@ -339,7 +333,7 @@ describe("useTrueFoundryAgentMessages", () => {
 
         const { result } = renderHook(() =>
             useTrueFoundryAgentMessages({
-                client: mockClient,
+                server: mockServer,
                 sessionId: "session-1",
                 getTurnHeaders,
             }),
@@ -353,7 +347,8 @@ describe("useTrueFoundryAgentMessages", () => {
 
         expect(getTurnHeaders).toHaveBeenCalledTimes(1);
         expect(streamTurnContent).toHaveBeenCalledWith(
-            expect.anything(),
+            mockServer,
+            "session-1",
             expect.any(PeerThreadFoldState),
             {
                 userMessage: "first",
@@ -371,7 +366,8 @@ describe("useTrueFoundryAgentMessages", () => {
 
         expect(getTurnHeaders).toHaveBeenCalledTimes(2);
         expect(streamTurnContent).toHaveBeenLastCalledWith(
-            expect.anything(),
+            mockServer,
+            "session-1",
             expect.any(PeerThreadFoldState),
             { userMessage: "second" },
             expect.any(AbortSignal),
@@ -385,11 +381,11 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
 
         await waitFor(() => expect(result.current.isLoading).toBe(false));
-        expect(loadSessionSnapshot).toHaveBeenCalledWith(mockClient, "session-1");
+        expect(loadSessionSnapshot).toHaveBeenCalledWith(mockServer, "session-1");
         expect(result.current.messages).toHaveLength(1);
         expect(result.current.messages[0]?.role).toBe("user");
     });
@@ -412,7 +408,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
 
         await waitFor(() => expect(result.current.isRunning).toBe(false));
@@ -455,7 +451,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
 
         await waitFor(() => expect(result.current.isLoading).toBe(false));
@@ -474,7 +470,7 @@ describe("useTrueFoundryAgentMessages", () => {
 
     it("sendTurn appends a user message and streams the assistant reply", async () => {
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.isLoading).toBe(false));
 
@@ -526,12 +522,6 @@ describe("useTrueFoundryAgentMessages", () => {
             }),
             fold,
         });
-        vi.mocked(getSession).mockResolvedValue({
-            cancel: vi.fn().mockResolvedValue(undefined),
-            listTurns: vi.fn(async function* () {}),
-            listEvents: vi.fn(async function* () {}),
-        } as never);
-
         let releaseStream: (() => void) | undefined;
         vi.mocked(streamTurnContent).mockReturnValue(
             (async function* () {
@@ -545,7 +535,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.isLoading).toBe(false));
         expect(result.current.messages.map((m) => m.role)).toEqual([
@@ -635,7 +625,7 @@ describe("useTrueFoundryAgentMessages", () => {
             );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.isLoading).toBe(false));
 
@@ -680,7 +670,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.isLoading).toBe(false));
 
@@ -700,7 +690,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.messages).toHaveLength(1));
 
@@ -718,7 +708,8 @@ describe("useTrueFoundryAgentMessages", () => {
         });
 
         expect(streamTurnContent).toHaveBeenCalledWith(
-            expect.anything(),
+            mockServer,
+            "session-1",
             expect.any(PeerThreadFoldState),
             {
                 inputs: [
@@ -743,7 +734,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.messages).toHaveLength(1));
 
@@ -777,7 +768,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.messages).toHaveLength(1));
 
@@ -799,7 +790,8 @@ describe("useTrueFoundryAgentMessages", () => {
 
         await waitFor(() => expect(streamTurnContent).toHaveBeenCalled());
         expect(streamTurnContent).toHaveBeenCalledWith(
-            expect.anything(),
+            mockServer,
+            "session-1",
             expect.any(PeerThreadFoldState),
             {
                 inputs: [
@@ -836,7 +828,7 @@ describe("useTrueFoundryAgentMessages", () => {
         );
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.messages).toHaveLength(2));
         expect(collectPendingToolResponses(result.current.messages)).toHaveLength(1);
@@ -858,13 +850,13 @@ describe("useTrueFoundryAgentMessages", () => {
     });
 
     describe("batched resume invariant", () => {
-        it("issues exactly one prepareTurn input batch across root and sub-agent threads", async () => {
+        it("issues exactly one prepareAndExecuteTurn input batch across root and sub-agent threads", async () => {
             vi.mocked(loadSessionSnapshot).mockResolvedValue(
                 snapshotWithAssistantMessage(assistantMessageWithMultiThreadPendingActions()),
             );
 
             const { result } = renderHook(() =>
-                useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+                useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
             );
             await waitFor(() => expect(result.current.messages).toHaveLength(1));
 
@@ -885,7 +877,8 @@ describe("useTrueFoundryAgentMessages", () => {
 
             await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(1));
             expect(streamTurnContent).toHaveBeenCalledWith(
-                expect.anything(),
+                mockServer,
+                "session-1",
                 expect.any(PeerThreadFoldState),
                 {
                     inputs: [
@@ -914,7 +907,7 @@ describe("useTrueFoundryAgentMessages", () => {
             );
 
             const { result } = renderHook(() =>
-                useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+                useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
             );
             await waitFor(() => expect(result.current.messages).toHaveLength(1));
 
@@ -931,7 +924,7 @@ describe("useTrueFoundryAgentMessages", () => {
         });
     });
 
-    it("cancel drains the stream gracefully and calls session.cancel", async () => {
+    it("cancel drains the stream gracefully and calls cancelSession", async () => {
         let resolveStream: (() => void) | undefined;
         vi.mocked(streamTurnContent).mockReturnValue(
             (async function* () {
@@ -941,15 +934,14 @@ describe("useTrueFoundryAgentMessages", () => {
                 });
             })(),
         );
-        // session.cancel() makes the backend close the SSE stream gracefully,
+        // cancelSession makes the backend close the SSE stream gracefully,
         // which ends the active iterator on its own.
-        const cancel = vi.fn().mockImplementation(async () => {
+        vi.mocked(mockServer.cancelSession).mockImplementation(async () => {
             resolveStream?.();
         });
-        vi.mocked(getSession).mockResolvedValue({ cancel } as never);
 
         const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+            useTrueFoundryAgentMessages({ server: mockServer, sessionId: "session-1" }),
         );
         await waitFor(() => expect(result.current.isLoading).toBe(false));
 
@@ -964,34 +956,10 @@ describe("useTrueFoundryAgentMessages", () => {
             await sendPromise;
         });
 
-        expect(cancel).toHaveBeenCalled();
+        expect(mockServer.cancelSession).toHaveBeenCalledWith({ sessionId: "session-1" });
         expect(result.current.isRunning).toBe(false);
         // No reconcile is triggered by cancel; the session was only loaded once
         // on mount and reconciles against the event log on the next page load.
         expect(loadSessionSnapshot).toHaveBeenCalledTimes(1);
     });
-
-    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 privateClient = {} as PrivateAgentSessionClient;
-
-        const { result } = renderHook(() =>
-            useTrueFoundryAgentMessages({
-                client: mockClient,
-                sessionId: "draft-session-1",
-                privateClient,
-            }),
-        );
-        await waitFor(() => expect(result.current.isLoading).toBe(false));
-
-        await act(async () => {
-            await result.current.cancel();
-        });
-
-        expect(cancel).toHaveBeenCalled();
-        expect(getSession).toHaveBeenCalledWith(mockClient, "draft-session-1", {
-            privateClient,
-        });
-    });
 });

package/src/useTrueFoundryAgentMessages.ts

@@ -3,15 +3,14 @@
 import { generateId } from "@assistant-ui/core";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import type {
-    AgentSessionClient,
     McpAuthRequiredEvent,
     ToolApprovalRequiredEvent,
     ToolResponseRequiredEvent,
     Turn,
     TurnInputItem,
     TurnStateDone,
-} from "truefoundry-gateway-sdk/agents";
-import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
+} from "./server/index.js";
+import type { AgentChatServer } from "./server/types.js";
 
 import { ROOT_THREAD_ID } from "./constants.js";
 import {
@@ -36,7 +35,6 @@ import {
     messageHasPendingRequiredActions,
     type RequiredActionInput,
 } from "./requiredActionInputs.js";
-import { getSession, type GetSessionOptions } from "./sessions.js";
 import {
     createEmptySessionSnapshot,
     replaceSessionSnapshot,
@@ -56,7 +54,7 @@ import {
 import type { TurnStreamUpdate } from "./turnStreamUpdate.js";
 
 export type UseTrueFoundryAgentMessagesOptions = {
-    client: AgentSessionClient;
+    server: AgentChatServer;
     sessionId: string | undefined;
     /** When true the thread is the currently selected (main) thread. */
     isMain?: boolean | undefined;
@@ -68,8 +66,6 @@ export type UseTrueFoundryAgentMessagesOptions = {
     }>;
     /** Maps a thread `remoteId` to the gateway session id used for turns. */
     resolveConversationSessionId?: (remoteId: string) => Promise<string>;
-    /** 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`.
@@ -287,22 +283,15 @@ function resolveTurnInput(
 }
 
 export function useTrueFoundryAgentMessages({
-    client,
+    server,
     sessionId,
     isMain,
     listEventsConcurrency,
     onError,
     initializeSession,
     resolveConversationSessionId,
-    privateClient,
     getTurnHeaders,
 }: UseTrueFoundryAgentMessagesOptions) {
-    const sessionOptions = useMemo<GetSessionOptions | undefined>(
-        () =>
-            privateClient != null ? { privateClient } : undefined,
-        [privateClient],
-    );
-
     const [snapshot, setSnapshot] = useState<SessionSnapshot>(createEmptySessionSnapshot);
     const [isRunning, setIsRunning] = useState(false);
     const [isLoading, setIsLoading] = useState(false);
@@ -508,9 +497,8 @@ export function useTrueFoundryAgentMessages({
                 resolveConversationSessionIdRef.current,
             );
             const loadedSnapshot = await loadSessionSnapshot(
-                client,
+                server,
                 conversationSessionId,
-                sessionOptions,
                 (snap) => {
                     if (generation === loadGenerationRef.current) {
                         setSnapshot(snap);
@@ -542,7 +530,9 @@ export function useTrueFoundryAgentMessages({
                 void runStream(
                     (signal) =>
                         resumeTurnStream(
-                            turn,
+                            server,
+                            conversationSessionId,
+                            turn.id,
                             loadedSnapshot.fold,
                             signal,
                             undefined,
@@ -562,7 +552,7 @@ export function useTrueFoundryAgentMessages({
                 setIsLoading(false);
             }
         }
-    }, [client, runStream, sessionId, sessionOptions, loadRetryTrigger, isMain]);
+    }, [server, runStream, sessionId, loadRetryTrigger, isMain]);
 
     useEffect(() => {
         void load().catch(() => undefined);
@@ -584,7 +574,6 @@ export function useTrueFoundryAgentMessages({
                 activeSessionId,
                 resolveConversationSessionIdRef.current,
             );
-            const session = await getSession(client, conversationSessionId, sessionOptions);
             const turnHeaders = await getTurnHeadersRef.current?.();
             const streamHeaders =
                 turnHeaders != null ? { headers: turnHeaders } : {};
@@ -596,7 +585,7 @@ export function useTrueFoundryAgentMessages({
                 isContinuation && continuationTurnId != null
                     ? continuationTurnId
                     : generateId();
-            // First turns must send previousTurnId: null.
+            // First turns must send previousTurnId: "none".
             const isFirstTurnInSession =
                 "userMessage" in options &&
                 options.previousTurnId === undefined &&
@@ -673,7 +662,8 @@ export function useTrueFoundryAgentMessages({
                 (signal) => {
                     if ("inputs" in options) {
                         return streamTurnContent(
-                            session,
+                            server,
+                            conversationSessionId,
                             snapshotRef.current.fold,
                             { inputs: options.inputs, ...streamHeaders },
                             signal,
@@ -682,7 +672,8 @@ export function useTrueFoundryAgentMessages({
                     }
                     if ("resumeMcpAuth" in options) {
                         return streamTurnContent(
-                            session,
+                            server,
+                            conversationSessionId,
                             snapshotRef.current.fold,
                             { resumeMcpAuth: true, ...streamHeaders },
                             signal,
@@ -690,14 +681,15 @@ export function useTrueFoundryAgentMessages({
                         );
                     }
                     return streamTurnContent(
-                        session,
+                        server,
+                        conversationSessionId,
                         snapshotRef.current.fold,
                         {
                             userMessage: options.userMessage,
                             ...(options.previousTurnId !== undefined
-                                ? { previousTurnId: options.previousTurnId }
+                                ? { previousTurnId: options.previousTurnId ?? "none" }
                                 : isFirstTurnInSession
-                                  ? { previousTurnId: null }
+                                  ? { previousTurnId: "none" }
                                   : {}),
                             ...streamHeaders,
                         },
@@ -726,7 +718,7 @@ export function useTrueFoundryAgentMessages({
                 isContinuation,
             );
         },
-        [client, runStream, sessionId, sessionOptions],
+        [server, runStream, sessionId],
     );
 
     const cancel = useCallback(async () => {
@@ -738,17 +730,16 @@ export function useTrueFoundryAgentMessages({
             sessionId,
             resolveConversationSessionIdRef.current,
         );
-        const session = await getSession(client, conversationSessionId, sessionOptions);
         // Request cancellation but keep consuming the stream. After cancel(),
         // the backend gracefully closes the SSE stream: it emits a terminal
         // turn.done event and then ends the stream, which lets the active run
         // drain to completion on its own instead of being torn down mid-flight.
-        await session.cancel().catch(() => undefined);
+        await server.cancelSession({ sessionId: conversationSessionId }).catch(() => undefined);
         // Wait for the in-flight stream to finish draining. No explicit
         // reconcile is needed here — the cancelled turn is terminal and local
         // state reconciles against the event log on the next session load.
         await activeRunRef.current?.catch(() => undefined);
-    }, [client, sessionId, sessionOptions]);
+    }, [server, sessionId]);
 
     const isRunningRef = useRef(isRunning);
     isRunningRef.current = isRunning;
@@ -817,7 +808,9 @@ export function useTrueFoundryAgentMessages({
         await runStream(
             (signal) =>
                 resumeTurnStream(
-                    turn,
+                    server,
+                    turn.sessionId,
+                    turn.id,
                     snapshotRef.current.fold,
                     signal,
                     undefined,
@@ -826,7 +819,7 @@ export function useTrueFoundryAgentMessages({
             { current: turn.id },
             true,
         );
-    }, [runStream]);
+    }, [runStream, server]);
 
     const branchFromTurn = useCallback(
         async (turnId: string, userMessage: UserMessageContent) => {
@@ -844,13 +837,14 @@ export function useTrueFoundryAgentMessages({
                 activeSessionId,
                 resolveConversationSessionIdRef.current,
             );
-            const session = await getSession(client, conversationSessionId, sessionOptions);
             const previousTurnId = await resolveGatewayBranchPreviousTurnIdForTurn(
-                session,
+                server,
+                conversationSessionId,
                 turnId,
             );
             const rewound = await buildSnapshotBeforeTurn(
-                session,
+                server,
+                conversationSessionId,
                 turnId,
                 listEventsConcurrency,
             );
@@ -868,11 +862,10 @@ export function useTrueFoundryAgentMessages({
         },
         [
             cancel,
-            client,
+            server,
             listEventsConcurrency,
             sendTurn,
             sessionId,
-            sessionOptions,
         ],
     );
 
@@ -936,13 +929,9 @@ export function useTrueFoundryAgentMessages({
                     sessionId,
                     resolveConversationSessionIdRef.current,
                 );
-                const session = await getSession(
-                    client,
-                    conversationSessionId,
-                    sessionOptions,
-                );
                 const next = await prependOlderSessionHistory(
-                    session,
+                    server,
+                    conversationSessionId,
                     snapshotRef.current,
                 );
                 if (generation !== loadGenerationRef.current) {
@@ -964,7 +953,7 @@ export function useTrueFoundryAgentMessages({
 
         loadOlderInflightRef.current = run;
         return run;
-    }, [client, isMain, sessionId, sessionOptions]);
+    }, [server, isMain, sessionId]);
 
     return {
         messages,
@@ -986,4 +975,3 @@ export function useTrueFoundryAgentMessages({
 }
 
 export { findPausedAssistantMessage, MCP_AUTH_RESUME_RUN_CUSTOM_KEY };
-