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

package/src/streamTurn.test.ts

@@ -1,15 +1,12 @@
 import { describe, expect, it, vi } from "vitest";
-import type {
-    AgentSession,
-    Turn,
-    TurnStreamData,
-} from "truefoundry-gateway-sdk/agents";
+import type { AgentChatServer, TurnStreamData } from "./server/index.js";
 
 import { ROOT_THREAD_ID } from "./constants.js";
 import { PeerThreadFoldState } from "./foldPeerThreads.js";
 import { resumeTurnStream, streamTurnContent } from "./streamTurn.js";
 
 const createdAt = new Date().toISOString();
+const SESSION_ID = "session-1";
 
 function streamData(
     sequenceNumber: number,
@@ -18,6 +15,10 @@ function streamData(
     return { sequenceNumber, event: event as TurnStreamData["event"] };
 }
 
+function mockServer(partial: Record<string, unknown>): AgentChatServer {
+    return partial as unknown as AgentChatServer;
+}
+
 async function collectUpdates(
     generator: AsyncGenerator<{ content: unknown[] }>,
 ): Promise<{ content: unknown[] }[]> {
@@ -32,42 +33,42 @@ describe("streamTurn", () => {
     describe("streamTurnContent", () => {
         it("prepares a user turn and yields folded stream updates", async () => {
             const foldState = new PeerThreadFoldState();
-            const execute = vi.fn(() =>
-                (async function* () {
-                    yield streamData(1, {
-                        type: "model.message",
-                        createdAt,
-                        id: "m1",
-                        threadId: ROOT_THREAD_ID,
-                        content: "hello from stream",
-                    });
-                })(),
-            );
-            const prepareTurn = vi.fn(() => ({ execute }));
-            const session = {
-                prepareTurn,
-                cancel: vi.fn().mockResolvedValue(undefined),
-            } as unknown as AgentSession;
+            const prepareAndExecuteTurn = vi.fn(async function* () {
+                yield streamData(1, {
+                    type: "model.message",
+                    createdAt,
+                    id: "m1",
+                    threadId: ROOT_THREAD_ID,
+                    content: "hello from stream",
+                });
+            });
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
 
             const updates = await collectUpdates(
                 streamTurnContent(
-                    session,
+                    server,
+                    SESSION_ID,
                     foldState,
                     { userMessage: "hello" },
                     new AbortController().signal,
                 ),
             );
 
-            expect(prepareTurn).toHaveBeenCalledWith({
+            expect(prepareAndExecuteTurn).toHaveBeenCalledWith({
+                sessionId: SESSION_ID,
                 input: [{ type: "user.message", content: "hello" }],
                 previousTurnId: "auto",
+                abortSignal: expect.any(AbortSignal),
             });
             expect(updates).toEqual([
                 { content: [{ type: "text", text: "hello from stream" }] },
             ]);
         });
 
-        it("passes required-action inputs through prepareTurn", async () => {
+        it("passes required-action inputs through prepareAndExecuteTurn", async () => {
             const inputs = [
                 {
                     type: "user.tool_approval" as const,
@@ -82,131 +83,139 @@ describe("streamTurn", () => {
                     content: "A",
                 },
             ];
-            const execute = vi.fn(() => (async function* () {})());
-            const prepareTurn = vi.fn(() => ({ execute }));
-            const session = {
-                prepareTurn,
-                cancel: vi.fn().mockResolvedValue(undefined),
-            } as unknown as AgentSession;
+            const prepareAndExecuteTurn = vi.fn(async function* () {});
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
 
             await collectUpdates(
                 streamTurnContent(
-                    session,
+                    server,
+                    SESSION_ID,
                     new PeerThreadFoldState(),
                     { inputs },
                     new AbortController().signal,
                 ),
             );
 
-            expect(prepareTurn).toHaveBeenCalledWith({
+            expect(prepareAndExecuteTurn).toHaveBeenCalledWith({
+                sessionId: SESSION_ID,
                 input: inputs,
                 previousTurnId: "auto",
+                abortSignal: expect.any(AbortSignal),
             });
         });
 
         it("uses empty input when resuming after MCP auth", 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 prepareAndExecuteTurn = vi.fn(async function* () {});
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
 
             await collectUpdates(
                 streamTurnContent(
-                    session,
+                    server,
+                    SESSION_ID,
                     new PeerThreadFoldState(),
                     { resumeMcpAuth: true },
                     new AbortController().signal,
                 ),
             );
 
-            expect(prepareTurn).toHaveBeenCalledWith({
+            expect(prepareAndExecuteTurn).toHaveBeenCalledWith({
+                sessionId: SESSION_ID,
                 input: [],
                 previousTurnId: "auto",
+                abortSignal: expect.any(AbortSignal),
             });
         });
 
         it("forwards an explicit previousTurnId when branching", 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 prepareAndExecuteTurn = vi.fn(async function* () {});
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
 
             await collectUpdates(
                 streamTurnContent(
-                    session,
+                    server,
+                    SESSION_ID,
                     new PeerThreadFoldState(),
                     { userMessage: "edited", previousTurnId: "turn-a" },
                     new AbortController().signal,
                 ),
             );
 
-            expect(prepareTurn).toHaveBeenCalledWith({
+            expect(prepareAndExecuteTurn).toHaveBeenCalledWith({
+                sessionId: SESSION_ID,
                 input: [{ type: "user.message", content: "edited" }],
                 previousTurnId: "turn-a",
+                abortSignal: expect.any(AbortSignal),
             });
         });
 
         it("forwards previousTurnId \"none\" when branching from root", 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 prepareAndExecuteTurn = vi.fn(async function* () {});
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
 
             await collectUpdates(
                 streamTurnContent(
-                    session,
+                    server,
+                    SESSION_ID,
                     new PeerThreadFoldState(),
                     { userMessage: "first", previousTurnId: "none" },
                     new AbortController().signal,
                 ),
             );
 
-            expect(prepareTurn).toHaveBeenCalledWith({
+            expect(prepareAndExecuteTurn).toHaveBeenCalledWith({
+                sessionId: SESSION_ID,
                 input: [{ type: "user.message", content: "first" }],
                 previousTurnId: "none",
+                abortSignal: expect.any(AbortSignal),
             });
         });
 
         it("returns early and cancels the session when already aborted", async () => {
-            const execute = vi.fn(() => (async function* () {})());
-            const prepareTurn = vi.fn(() => ({ execute }));
-            const cancel = vi.fn().mockResolvedValue(undefined);
-            const session = { prepareTurn, cancel } as unknown as AgentSession;
+            const prepareAndExecuteTurn = vi.fn(async function* () {});
+            const cancelSession = vi.fn().mockResolvedValue(undefined);
+            const server = mockServer({ prepareAndExecuteTurn, cancelSession });
             const abortController = new AbortController();
             abortController.abort();
 
             const updates = await collectUpdates(
                 streamTurnContent(
-                    session,
+                    server,
+                    SESSION_ID,
                     new PeerThreadFoldState(),
                     { userMessage: "hello" },
                     abortController.signal,
                 ),
             );
 
-            expect(cancel).toHaveBeenCalled();
-            expect(execute).not.toHaveBeenCalled();
+            expect(cancelSession).toHaveBeenCalledWith({ sessionId: SESSION_ID });
+            expect(prepareAndExecuteTurn).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;
+        it("forwards headers to prepareAndExecuteTurn", async () => {
+            const prepareAndExecuteTurn = vi.fn(async function* () {});
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
             const abortSignal = new AbortController().signal;
 
             await collectUpdates(
                 streamTurnContent(
-                    session,
+                    server,
+                    SESSION_ID,
                     new PeerThreadFoldState(),
                     {
                         userMessage: "hello",
@@ -218,54 +227,50 @@ describe("streamTurn", () => {
                 ),
             );
 
-            expect(execute).toHaveBeenCalledWith(
-                { stream: true },
-                {
-                    abortSignal,
-                    headers: {
-                        "x-tfy-session-last-updated-at": "2026-06-30T10:00:00.000Z",
-                    },
+            expect(prepareAndExecuteTurn).toHaveBeenCalledWith({
+                sessionId: SESSION_ID,
+                input: [{ type: "user.message", content: "hello" }],
+                previousTurnId: "auto",
+                abortSignal,
+                headers: {
+                    "x-tfy-session-last-updated-at": "2026-06-30T10:00:00.000Z",
                 },
-            );
+            });
         });
 
         it("notifies gateway turn id when turn.done errors with no content yields", async () => {
             const gatewayTurnId = "01ky6mqzmczwt6ssyd5r02gjjc";
-            const turn = {
-                id: undefined as string | undefined,
-                execute: vi.fn(() =>
-                    (async function* () {
-                        // SDK sets turn.id when turn.created is observed.
-                        turn.id = gatewayTurnId;
-                        yield streamData(1, {
-                            type: "turn.created",
-                            createdAt,
-                            input: [{ type: "user.message", content: "hello" }],
-                        });
-                        yield streamData(2, {
-                            type: "turn.done",
-                            createdAt,
-                            state: {
-                                status: "error",
-                                message:
-                                    "Publisher Model is not servable in region us-central1.",
-                                completedAt: createdAt,
-                            },
-                        });
-                    })(),
-                ),
-            };
-            const prepareTurn = vi.fn(() => turn);
-            const session = {
-                prepareTurn,
-                cancel: vi.fn().mockResolvedValue(undefined),
-            } as unknown as AgentSession;
+            const prepareAndExecuteTurn = vi.fn(async function* () {
+                yield streamData(1, {
+                    type: "turn.created",
+                    createdAt,
+                    id: "created-1",
+                    turnId: gatewayTurnId,
+                    input: [{ type: "user.message", content: "hello" }],
+                });
+                yield streamData(2, {
+                    type: "turn.done",
+                    createdAt,
+                    id: "done-1",
+                    state: {
+                        status: "error",
+                        message:
+                            "Publisher Model is not servable in region us-central1.",
+                        completedAt: createdAt,
+                    },
+                });
+            });
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
             const onTurnIdAvailable = vi.fn();
 
             await expect(
                 collectUpdates(
                     streamTurnContent(
-                        session,
+                        server,
+                        SESSION_ID,
                         new PeerThreadFoldState(),
                         { userMessage: "hello" },
                         new AbortController().signal,
@@ -279,31 +284,30 @@ describe("streamTurn", () => {
             expect(onTurnIdAvailable).toHaveBeenCalledWith(gatewayTurnId);
         });
 
-        it("does not notify when an error stream never assigns turn.id", async () => {
-            const execute = vi.fn(() =>
-                (async function* () {
-                    yield streamData(1, {
-                        type: "turn.done",
-                        createdAt,
-                        state: {
-                            status: "error",
-                            completedAt: createdAt,
-                            message: "boom",
-                        },
-                    });
-                })(),
-            );
-            const prepareTurn = vi.fn(() => ({ id: undefined, execute }));
-            const session = {
-                prepareTurn,
-                cancel: vi.fn().mockResolvedValue(undefined),
-            } as unknown as AgentSession;
+        it("does not notify when an error stream never emits turn.created", async () => {
+            const prepareAndExecuteTurn = vi.fn(async function* () {
+                yield streamData(1, {
+                    type: "turn.done",
+                    createdAt,
+                    id: "done-1",
+                    state: {
+                        status: "error",
+                        completedAt: createdAt,
+                        message: "boom",
+                    },
+                });
+            });
+            const server = mockServer({
+                prepareAndExecuteTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
             const onTurnIdAvailable = vi.fn();
 
             await expect(
                 collectUpdates(
                     streamTurnContent(
-                        session,
+                        server,
+                        SESSION_ID,
                         new PeerThreadFoldState(),
                         { userMessage: "hello" },
                         new AbortController().signal,
@@ -320,45 +324,58 @@ describe("streamTurn", () => {
     describe("resumeTurnStream", () => {
         it("reconnects with afterSequenceNumber and yields updates", async () => {
             const foldState = new PeerThreadFoldState();
-            const stream = vi.fn(() =>
-                (async function* () {
-                    yield streamData(2, {
-                        type: "model.message",
-                        createdAt,
-                        id: "m2",
-                        threadId: ROOT_THREAD_ID,
-                        content: "resumed",
-                    });
-                })(),
-            );
-            const turn = {
-                stream,
-                session: { cancel: vi.fn().mockResolvedValue(undefined) },
-            } as unknown as Turn;
+            const subscribeToTurn = vi.fn(async function* () {
+                yield streamData(2, {
+                    type: "model.message",
+                    createdAt,
+                    id: "m2",
+                    threadId: ROOT_THREAD_ID,
+                    content: "resumed",
+                });
+            });
+            const server = mockServer({
+                subscribeToTurn,
+                cancelSession: vi.fn().mockResolvedValue(undefined),
+            });
 
             const updates = await collectUpdates(
-                resumeTurnStream(turn, foldState, new AbortController().signal, 1),
+                resumeTurnStream(
+                    server,
+                    SESSION_ID,
+                    "turn-1",
+                    foldState,
+                    new AbortController().signal,
+                    1,
+                ),
             );
-            expect(stream).toHaveBeenCalledWith({ afterSequenceNumber: 1 }, expect.any(Object));
+            expect(subscribeToTurn).toHaveBeenCalledWith({
+                sessionId: SESSION_ID,
+                turnId: "turn-1",
+                afterSequenceNumber: 1,
+                abortSignal: expect.any(AbortSignal),
+            });
             expect(updates).toEqual([{ content: [{ type: "text", text: "resumed" }] }]);
         });
 
         it("returns early when aborted before streaming starts", async () => {
-            const stream = vi.fn(() => (async function* () {})());
-            const cancel = vi.fn().mockResolvedValue(undefined);
-            const turn = {
-                stream,
-                session: { cancel },
-            } as unknown as Turn;
+            const subscribeToTurn = vi.fn(async function* () {});
+            const cancelSession = vi.fn().mockResolvedValue(undefined);
+            const server = mockServer({ subscribeToTurn, cancelSession });
             const abortController = new AbortController();
             abortController.abort();
 
             const updates = await collectUpdates(
-                resumeTurnStream(turn, new PeerThreadFoldState(), abortController.signal),
+                resumeTurnStream(
+                    server,
+                    SESSION_ID,
+                    "turn-1",
+                    new PeerThreadFoldState(),
+                    abortController.signal,
+                ),
             );
 
-            expect(cancel).toHaveBeenCalled();
-            expect(stream).not.toHaveBeenCalled();
+            expect(cancelSession).toHaveBeenCalledWith({ sessionId: SESSION_ID });
+            expect(subscribeToTurn).not.toHaveBeenCalled();
             expect(updates).toEqual([]);
         });
     });