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

package/src/requiredActionsFromActiveUpdate.test.ts

@@ -0,0 +1,95 @@
+import { describe, expect, it } from "vitest";
+
+import { ROOT_THREAD_ID } from "./constants.js";
+import { toolApprovalStatus, TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY } from "./toolApproval.js";
+import { toolResponseStatus, TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY } from "./toolResponse.js";
+import type { TurnStreamUpdate } from "./turnStreamUpdate.js";
+import { requiredActionsFromActiveUpdate } from "./useTrueFoundryAgentMessages.js";
+
+/**
+ * Regression coverage for the live-stream pause path: when a turn pauses at a
+ * tool approval or an `ask_user_question` (`tool.response_required`),
+ * `commitActiveStream` fabricates a synthetic `TurnStateDone`. Its
+ * `requiredActions` are reconstructed here from the paused update. If they were
+ * dropped (as they were before the fix), the committed assistant message would
+ * lose its `requires-action` status and the resume turn would never be sent.
+ */
+describe("requiredActionsFromActiveUpdate", () => {
+    it("reconstructs a tool.response_required pause (ask_user_question)", () => {
+        const update: TurnStreamUpdate = {
+            content: [],
+            status: toolResponseStatus(),
+            metadata: {
+                custom: { [TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY]: ROOT_THREAD_ID },
+            },
+        };
+
+        expect(requiredActionsFromActiveUpdate(update)).toEqual([
+            expect.objectContaining({
+                type: "tool.response_required",
+                threadId: ROOT_THREAD_ID,
+                toolCalls: [],
+            }),
+        ]);
+    });
+
+    it("reconstructs a tool.approval_required pause", () => {
+        const update: TurnStreamUpdate = {
+            content: [],
+            status: toolApprovalStatus(),
+            metadata: {
+                custom: { [TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY]: "sub-agent-1" },
+            },
+        };
+
+        expect(requiredActionsFromActiveUpdate(update)).toEqual([
+            expect.objectContaining({
+                type: "tool.approval_required",
+                threadId: "sub-agent-1",
+                toolCalls: [],
+            }),
+        ]);
+    });
+
+    it("reconstructs both approval and response pauses in one turn", () => {
+        const update: TurnStreamUpdate = {
+            content: [],
+            status: toolApprovalStatus(),
+            metadata: {
+                custom: {
+                    [TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY]: ROOT_THREAD_ID,
+                    [TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY]: "sub-agent-1",
+                },
+            },
+        };
+
+        const actions = requiredActionsFromActiveUpdate(update);
+        expect(actions.map((action) => action.type)).toEqual([
+            "tool.approval_required",
+            "tool.response_required",
+        ]);
+    });
+
+    it("reconstructs an mcp.auth_required pause", () => {
+        const mcpServers = [{ id: "gh", name: "GitHub", authUrl: "https://auth" }];
+        const update: TurnStreamUpdate = {
+            content: [],
+            metadata: { custom: { pendingMcpAuth: true, mcpServers } },
+        };
+
+        expect(requiredActionsFromActiveUpdate(update)).toEqual([
+            expect.objectContaining({ type: "mcp.auth_required", mcpServers }),
+        ]);
+    });
+
+    it("returns no actions for a completed (non-paused) update", () => {
+        const update: TurnStreamUpdate = {
+            content: [],
+            metadata: {
+                custom: { [TOOL_RESPONSE_THREAD_ID_CUSTOM_KEY]: ROOT_THREAD_ID },
+            },
+        };
+
+        expect(requiredActionsFromActiveUpdate(update)).toEqual([]);
+    });
+});

package/src/sessionListStartTimestamp.ts

@@ -0,0 +1,6 @@
+/** ISO timestamp one year before now — used as `listSessions` `startTimestamp`. */
+export function sessionListStartTimestamp(): string {
+    const start = new Date();
+    start.setFullYear(start.getFullYear() - 1);
+    return start.toISOString();
+}

package/src/sessionSnapshot.ts

@@ -0,0 +1,107 @@
+import type { McpAuthRequiredEvent, Turn, TurnInputItem } from "truefoundry-gateway-sdk/agents";
+
+import { extractTurnUserText } from "./extractTurnUserText.js";
+import { PeerThreadFoldState } from "./foldPeerThreads.js";
+import type { StoredApprovalDecision } from "./toolApproval.js";
+import type { StoredToolResponse } from "./toolResponse.js";
+import type { TurnStreamUpdate } from "./turnStreamUpdate.js";
+
+/** Turn metadata retained for cross-turn projection and subsequent required-action replay. */
+export type SessionTurnRecord = Pick<
+    Turn,
+    "id" | "createdAt" | "state" | "input"
+> & {
+    /** Denormalized from `input` for user/assistant interleaving during projection. */
+    userText?: string;
+    /** Root-thread `model.message` ids ingested with this turn (for per-group projection). */
+    rootModelMessageIds?: readonly string[];
+    /** sandboxId observed via a `sandbox.created` event during this turn (session-scoped, latest wins). */
+    sandboxId?: string;
+};
+
+export type RequiredActionsOverlay = {
+    approvals: Map<string, StoredApprovalDecision>;
+    toolResponses: Map<string, StoredToolResponse>;
+};
+
+export type ActiveStreamState = {
+    turnId: string;
+    update: TurnStreamUpdate;
+    isContinuation: boolean;
+    streamComplete?: boolean;
+};
+
+export type PendingUserMessage = {
+    turnId: string;
+    /** Gateway user.message content (text-only string or text/file parts). */
+    content: Extract<TurnInputItem, { type: "user.message" }>["content"];
+    createdAt: Date;
+};
+
+export type SessionSnapshot = {
+    fold: PeerThreadFoldState;
+    turns: SessionTurnRecord[];
+    pendingMcpAuth?: McpAuthRequiredEvent;
+    pendingUser?: PendingUserMessage;
+    activeStream?: ActiveStreamState;
+    /** Root `model.message` ids present before the active turn group started (streaming scope). */
+    groupRootBaseline?: readonly string[];
+    requiredActions: RequiredActionsOverlay;
+    runningTurn?: Turn;
+    unstable_resume?: boolean;
+};
+
+export type ProjectSessionMessagesOptions = {
+    getCreatedAt?: (messageId: string, fallback: Date) => Date;
+};
+
+export function emptyRequiredActionsOverlay(): RequiredActionsOverlay {
+    return {
+        approvals: new Map(),
+        toolResponses: new Map(),
+    };
+}
+
+export function createEmptySessionSnapshot(): SessionSnapshot {
+    return {
+        fold: new PeerThreadFoldState(),
+        turns: [],
+        requiredActions: emptyRequiredActionsOverlay(),
+    };
+}
+
+export function turnToSessionRecord(turn: Turn): SessionTurnRecord {
+    const userText = extractTurnUserText(turn.input);
+    return {
+        id: turn.id,
+        ...(userText ? { userText } : {}),
+        createdAt: turn.createdAt,
+        state: turn.state,
+        input: turn.input,
+    };
+}
+
+/** 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">> & {
+        requiredActions?: RequiredActionsOverlay;
+    },
+): SessionSnapshot {
+    return {
+        ...snapshot,
+        ...patch,
+        ...(patch.requiredActions != null
+            ? { requiredActions: patch.requiredActions }
+            : {}),
+    };
+}
+
+export function cloneRequiredActionsOverlay(
+    overlay: RequiredActionsOverlay,
+): RequiredActionsOverlay {
+    return {
+        approvals: new Map(overlay.approvals),
+        toolResponses: new Map(overlay.toolResponses),
+    };
+}

package/src/sessions.ts

@@ -0,0 +1,36 @@
+import type {
+    AgentSession,
+    AgentSessionClient,
+} from "truefoundry-gateway-sdk/agents";
+import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
+
+import { bindDraftAgentSession } from "./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;
+};
+
+/** `sessionId` is the assistant-ui thread `remoteId` from `RemoteThreadListAdapter.initialize`. */
+export function getSession(
+    client: AgentSessionClient,
+    sessionId: string,
+    options?: GetSessionOptions,
+): Promise<AgentSession> {
+    if (options?.draftGateway != null) {
+        return bindDraftAgentSession(client, options.draftGateway, sessionId);
+    }
+
+    let inflight = inflightBySessionId.get(sessionId);
+    if (inflight == null) {
+        inflight = client.getSession({ sessionId }).finally(() => {
+            if (inflightBySessionId.get(sessionId) === inflight) {
+                inflightBySessionId.delete(sessionId);
+            }
+        });
+        inflightBySessionId.set(sessionId, inflight);
+    }
+    return inflight;
+}

package/src/streamTurn.test.ts

@@ -0,0 +1,243 @@
+import { describe, expect, it, vi } from "vitest";
+import type {
+    AgentSession,
+    Turn,
+    TurnStreamData,
+} from "truefoundry-gateway-sdk/agents";
+
+import { ROOT_THREAD_ID } from "./constants.js";
+import { PeerThreadFoldState } from "./foldPeerThreads.js";
+import { resumeTurnStream, streamTurnContent } from "./streamTurn.js";
+
+const createdAt = new Date().toISOString();
+
+function streamData(
+    sequenceNumber: number,
+    event: TurnStreamData["event"],
+): TurnStreamData {
+    return { sequenceNumber, event };
+}
+
+async function collectUpdates(
+    generator: AsyncGenerator<{ content: unknown[] }>,
+): Promise<{ content: unknown[] }[]> {
+    const updates: { content: unknown[] }[] = [];
+    for await (const update of generator) {
+        updates.push(update);
+    }
+    return updates;
+}
+
+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 updates = await collectUpdates(
+                streamTurnContent(
+                    session,
+                    foldState,
+                    { userMessage: "hello" },
+                    new AbortController().signal,
+                ),
+            );
+
+            expect(prepareTurn).toHaveBeenCalledWith({
+                input: [{ type: "user.message", content: "hello" }],
+                previousTurnId: "auto",
+            });
+            expect(updates).toEqual([
+                { content: [{ type: "text", text: "hello from stream" }] },
+            ]);
+        });
+
+        it("passes required-action inputs through prepareTurn", async () => {
+            const inputs = [
+                {
+                    type: "user.tool_approval" as const,
+                    threadId: ROOT_THREAD_ID,
+                    toolCallId: "approval-1",
+                    approval: { status: "allow" as const },
+                },
+                {
+                    type: "user.tool_response" as const,
+                    threadId: ROOT_THREAD_ID,
+                    toolCallId: "question-1",
+                    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;
+
+            await collectUpdates(
+                streamTurnContent(
+                    session,
+                    new PeerThreadFoldState(),
+                    { inputs },
+                    new AbortController().signal,
+                ),
+            );
+
+            expect(prepareTurn).toHaveBeenCalledWith({
+                input: inputs,
+                previousTurnId: "auto",
+            });
+        });
+
+        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;
+
+            await collectUpdates(
+                streamTurnContent(
+                    session,
+                    new PeerThreadFoldState(),
+                    { resumeMcpAuth: true },
+                    new AbortController().signal,
+                ),
+            );
+
+            expect(prepareTurn).toHaveBeenCalledWith({
+                input: [],
+                previousTurnId: "auto",
+            });
+        });
+
+        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;
+
+            await collectUpdates(
+                streamTurnContent(
+                    session,
+                    new PeerThreadFoldState(),
+                    { userMessage: "edited", previousTurnId: "turn-a" },
+                    new AbortController().signal,
+                ),
+            );
+
+            expect(prepareTurn).toHaveBeenCalledWith({
+                input: [{ type: "user.message", content: "edited" }],
+                previousTurnId: "turn-a",
+            });
+        });
+
+        it("omits previousTurnId 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;
+
+            await collectUpdates(
+                streamTurnContent(
+                    session,
+                    new PeerThreadFoldState(),
+                    { userMessage: "first", previousTurnId: null },
+                    new AbortController().signal,
+                ),
+            );
+
+            expect(prepareTurn).toHaveBeenCalledWith({
+                input: [{ type: "user.message", content: "first" }],
+            });
+        });
+
+        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 abortController = new AbortController();
+            abortController.abort();
+
+            const updates = await collectUpdates(
+                streamTurnContent(
+                    session,
+                    new PeerThreadFoldState(),
+                    { userMessage: "hello" },
+                    abortController.signal,
+                ),
+            );
+
+            expect(cancel).toHaveBeenCalled();
+            expect(execute).not.toHaveBeenCalled();
+            expect(updates).toEqual([]);
+        });
+    });
+
+    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 updates = await collectUpdates(
+                resumeTurnStream(turn, foldState, new AbortController().signal, 1),
+            );
+            expect(stream).toHaveBeenCalledWith({ afterSequenceNumber: 1 }, expect.any(Object));
+            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 abortController = new AbortController();
+            abortController.abort();
+
+            const updates = await collectUpdates(
+                resumeTurnStream(turn, new PeerThreadFoldState(), abortController.signal),
+            );
+
+            expect(cancel).toHaveBeenCalled();
+            expect(stream).not.toHaveBeenCalled();
+            expect(updates).toEqual([]);
+        });
+    });
+});

package/src/streamTurn.ts

@@ -0,0 +1,119 @@
+import {
+    type AgentSession,
+    type Turn,
+    type TurnInputItem,
+} from "truefoundry-gateway-sdk/agents";
+
+import {
+    streamTurnEvents,
+    type UserMessageContent,
+} from "./convertTurnMessages.js";
+import { PeerThreadFoldState } from "./foldPeerThreads.js";
+import type { RequiredActionInput } from "./requiredActionInputs.js";
+import type { TurnStreamUpdate } from "./turnStreamUpdate.js";
+
+export type StreamTurnOptions = {
+    userMessage?: UserMessageContent;
+    resumeMcpAuth?: boolean;
+    inputs?: RequiredActionInput[];
+    /**
+     * Branch anchor for `prepareTurn`. Omit for `"auto"`. Pass `null` for a fresh
+     * root turn (no `previousTurnId` field).
+     */
+    previousTurnId?: string | null;
+};
+
+function buildTurnInput(options: StreamTurnOptions): TurnInputItem[] {
+    if (options.inputs != null) {
+        return options.inputs;
+    }
+    if (options.resumeMcpAuth === true) {
+        return [];
+    }
+    return [{ type: "user.message", content: options.userMessage ?? "" }];
+}
+
+function bindAbort(session: AgentSession, abortSignal: AbortSignal): () => void {
+    const onAbort = () => {
+        void session.cancel().catch(() => undefined);
+    };
+    if (abortSignal.aborted) {
+        onAbort();
+        return onAbort;
+    }
+    abortSignal.addEventListener("abort", onAbort, { once: true });
+    return onAbort;
+}
+
+export async function* streamTurnContent(
+    session: AgentSession,
+    foldState: PeerThreadFoldState,
+    options: StreamTurnOptions,
+    abortSignal: AbortSignal,
+    groupRootBaseline?: readonly string[],
+): AsyncGenerator<TurnStreamUpdate> {
+    const previousTurnId =
+        options.previousTurnId === null
+            ? undefined
+            : (options.previousTurnId ?? "auto");
+    const turn = session.prepareTurn({
+        input: buildTurnInput(options),
+        ...(previousTurnId != null ? { previousTurnId } : {}),
+    });
+
+    const onAbort = bindAbort(session, abortSignal);
+    if (abortSignal.aborted) {
+        return;
+    }
+
+    try {
+        yield* streamTurnEvents(
+            turn.execute({ stream: true }, { abortSignal }),
+            foldState,
+            groupRootBaseline,
+        );
+    } catch (error) {
+        if (error instanceof Error && error.name === "AbortError") {
+            return;
+        }
+        throw error;
+    } finally {
+        abortSignal.removeEventListener("abort", onAbort);
+    }
+}
+
+/** TODO: wire `afterSequenceNumber` from the last ingested stream event to skip replay on reconnect. */
+export async function* resumeTurnStream(
+    turn: Turn,
+    foldState: PeerThreadFoldState,
+    abortSignal: AbortSignal,
+    afterSequenceNumber?: number,
+    groupRootBaseline?: readonly string[],
+): AsyncGenerator<TurnStreamUpdate> {
+    const onAbort = () => {
+        void turn.session.cancel().catch(() => undefined);
+    };
+    if (abortSignal.aborted) {
+        onAbort();
+        return;
+    }
+    abortSignal.addEventListener("abort", onAbort, { once: true });
+
+    try {
+        yield* streamTurnEvents(
+            turn.stream(
+                afterSequenceNumber != null ? { afterSequenceNumber } : {},
+                { abortSignal },
+            ),
+            foldState,
+            groupRootBaseline,
+        );
+    } catch (error) {
+        if (error instanceof Error && error.name === "AbortError") {
+            return;
+        }
+        throw error;
+    } finally {
+        abortSignal.removeEventListener("abort", onAbort);
+    }
+}