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

package/src/useTrueFoundryAgentMessages.test.tsx

@@ -1,7 +1,7 @@
 // @vitest-environment jsdom
 import type { ThreadMessage } from "@assistant-ui/core";
 import { act, renderHook, waitFor } from "@testing-library/react";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+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";
 
@@ -288,6 +288,10 @@ describe("useTrueFoundryAgentMessages", () => {
         vi.mocked(resumeTurnStream).mockReturnValue(singleUpdateStream());
     });
 
+    afterEach(() => {
+        vi.unstubAllGlobals();
+    });
+
     it("clears messages when sessionId is undefined", async () => {
         const { result } = renderHook(() =>
             useTrueFoundryAgentMessages({ client: mockClient, sessionId: undefined }),
@@ -372,6 +376,52 @@ describe("useTrueFoundryAgentMessages", () => {
         );
     });
 
+    it("clears isLoading while a resumed turn is still streaming", async () => {
+        let releaseStream: (() => void) | undefined;
+        vi.mocked(resumeTurnStream).mockReturnValue(
+            (async function* () {
+                await new Promise<void>((resolve) => {
+                    releaseStream = resolve;
+                });
+                yield { content: [{ type: "text" as const, text: "streamed reply" }] };
+            })(),
+        );
+
+        const runningTurn = {
+            id: "turn-running",
+            input: [{ type: "user.message", content: "keep going" }],
+            createdAt: new Date().toISOString(),
+        } as Turn;
+        vi.mocked(loadSessionSnapshot).mockResolvedValue(
+            replaceSessionSnapshot(createEmptySessionSnapshot(), {
+                runningTurn,
+                unstable_resume: true,
+                pendingUser: {
+                    turnId: runningTurn.id,
+                    content: "keep going",
+                    createdAt: new Date(runningTurn.createdAt),
+                },
+            }),
+        );
+
+        const { result } = renderHook(() =>
+            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+        );
+
+        await waitFor(() => expect(result.current.isLoading).toBe(false));
+        expect(resumeTurnStream).toHaveBeenCalled();
+        await waitFor(() => expect(result.current.isRunning).toBe(true));
+        expect(result.current.messages[0]).toMatchObject({
+            role: "user",
+            content: [{ type: "text", text: "keep going" }],
+        });
+
+        await act(async () => {
+            releaseStream?.();
+        });
+        await waitFor(() => expect(result.current.isRunning).toBe(false));
+    });
+
     it("sendTurn appends a user message and streams the assistant reply", async () => {
         const { result } = renderHook(() =>
             useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
@@ -395,6 +445,180 @@ describe("useTrueFoundryAgentMessages", () => {
         });
     });
 
+    it("editFromTurn drops prior turns before showing the edited user message", async () => {
+        const createdAt = new Date().toISOString();
+        const fold = new PeerThreadFoldState();
+        ingestTurnEvent(fold, {
+            type: "model.message",
+            id: "model-1",
+            createdAt,
+            threadId: ROOT_THREAD_ID,
+            role: "assistant",
+            content: "How are you",
+        } as never);
+
+        vi.mocked(loadSessionSnapshot).mockResolvedValue({
+            ...replaceSessionSnapshot(createEmptySessionSnapshot(), {
+                turns: [
+                    {
+                        id: "turn-1",
+                        userText: "Hello",
+                        createdAt,
+                        state: {
+                            status: "done",
+                            requiredActions: [],
+                            completedAt: createdAt,
+                        },
+                        input: [{ type: "user.message", content: "Hello" }],
+                        rootModelMessageIds: ["model-1"],
+                    },
+                ],
+            }),
+            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* () {
+                yield {
+                    content: [{ type: "text" as const, text: "sunny" }],
+                };
+                await new Promise<void>((resolve) => {
+                    releaseStream = resolve;
+                });
+            })(),
+        );
+
+        const { result } = renderHook(() =>
+            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+        );
+        await waitFor(() => expect(result.current.isLoading).toBe(false));
+        expect(result.current.messages.map((m) => m.role)).toEqual([
+            "user",
+            "assistant",
+        ]);
+        expect(result.current.messages[0]).toMatchObject({
+            content: [{ type: "text", text: "Hello" }],
+        });
+
+        let editPromise: Promise<void>;
+        await act(async () => {
+            editPromise = result.current.editFromTurn(
+                "turn-1",
+                "what is the weather like?",
+            );
+            await Promise.resolve();
+        });
+
+        await waitFor(() => {
+            const texts = result.current.messages
+                .filter((m) => m.role === "user")
+                .map((m) =>
+                    m.content
+                        .filter((p): p is { type: "text"; text: string } => p.type === "text")
+                        .map((p) => p.text)
+                        .join(""),
+                );
+            expect(texts).toEqual(["what is the weather like?"]);
+        });
+        expect(
+            result.current.messages.some(
+                (m) =>
+                    m.role === "user" &&
+                    m.content.some((p) => p.type === "text" && p.text === "Hello"),
+            ),
+        ).toBe(false);
+        expect(
+            result.current.messages.some(
+                (m) =>
+                    m.role === "assistant" &&
+                    m.content.some(
+                        (p) => p.type === "text" && p.text.includes("How are you"),
+                    ),
+            ),
+        ).toBe(false);
+
+        await act(async () => {
+            releaseStream?.();
+            await editPromise!;
+        });
+    });
+
+    it("does not let a superseded stream complete the current stream", async () => {
+        let releaseFirstStream: (() => void) | undefined;
+        let releaseSecondStream: (() => void) | undefined;
+        let nextAnimationFrame = 1;
+        const animationFrames = new Map<number, FrameRequestCallback>();
+        vi.stubGlobal(
+            "requestAnimationFrame",
+            vi.fn((callback: FrameRequestCallback) => {
+                const frame = nextAnimationFrame++;
+                animationFrames.set(frame, callback);
+                return frame;
+            }),
+        );
+        vi.stubGlobal(
+            "cancelAnimationFrame",
+            vi.fn((frame: number) => animationFrames.delete(frame)),
+        );
+        vi.mocked(streamTurnContent)
+            .mockReturnValueOnce(
+                (async function* () {
+                    yield { content: [{ type: "text" as const, text: "first reply" }] };
+                    await new Promise<void>((resolve) => {
+                        releaseFirstStream = resolve;
+                    });
+                })(),
+            )
+            .mockReturnValueOnce(
+                (async function* () {
+                    yield { content: [{ type: "text" as const, text: "second reply" }] };
+                    await new Promise<void>((resolve) => {
+                        releaseSecondStream = resolve;
+                    });
+                })(),
+            );
+
+        const { result } = renderHook(() =>
+            useTrueFoundryAgentMessages({ client: mockClient, sessionId: "session-1" }),
+        );
+        await waitFor(() => expect(result.current.isLoading).toBe(false));
+
+        const firstSend = result.current.sendTurn({ userMessage: "first" });
+        await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(1));
+
+        const secondSend = result.current.sendTurn({ userMessage: "second" });
+        await waitFor(() => expect(streamTurnContent).toHaveBeenCalledTimes(2));
+
+        await act(async () => {
+            releaseFirstStream?.();
+            await firstSend;
+        });
+
+        expect(result.current.isRunning).toBe(true);
+
+        await act(async () => {
+            animationFrames.get(2)?.(performance.now());
+        });
+        expect(result.current.messages.at(-1)).toMatchObject({
+            role: "assistant",
+            content: [{ type: "text", text: "second reply" }],
+        });
+        expect(result.current.messages.at(-1)?.status).not.toMatchObject({
+            type: "complete",
+        });
+
+        await act(async () => {
+            releaseSecondStream?.();
+            await secondSend;
+        });
+    });
+
     it("carries a streamed sandboxId through commit so it survives after the stream completes", async () => {
         vi.mocked(streamTurnContent).mockReturnValue(
             (async function* () {