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

package/src/streamTurn.ts

@@ -1,9 +1,9 @@
 import type {
-    AgentSession,
-    Turn,
+    AgentChatServer,
+    PreviousTurnIdInput,
     TurnInputItem,
-} from "truefoundry-gateway-sdk/agents";
-import type { TruefoundryGatewayApi } from "truefoundry-gateway-sdk";
+} from "./server/types.js";
+import type { TurnStreamData } from "./server/events.js";
 
 import {
     streamTurnEvents,
@@ -18,11 +18,11 @@ export type StreamTurnOptions = {
     resumeMcpAuth?: boolean;
     inputs?: RequiredActionInput[];
     /**
-     * Branch anchor for `prepareTurn`. Omit for `"auto"`. Pass `"none"` for a fresh
-     * root turn — sent as `previous_turn_id: "none"` on the wire.
+     * Branch anchor for prepareAndExecuteTurn. Omit for `"auto"`. Pass `"none"` for a fresh
+     * root turn.
      */
-    previousTurnId?: TruefoundryGatewayApi.PreviousTurnIdInput;
-    /** Extra headers for the createTurn request (`execute` request options). */
+    previousTurnId?: PreviousTurnIdInput;
+    /** Extra headers for the turn request. */
     headers?: Record<string, string>;
 };
 
@@ -36,9 +36,13 @@ function buildTurnInput(options: StreamTurnOptions): TurnInputItem[] {
     return [{ type: "user.message", content: options.userMessage ?? "" }];
 }
 
-function bindAbort(session: AgentSession, abortSignal: AbortSignal): () => void {
+function bindAbort(
+    server: AgentChatServer,
+    sessionId: string,
+    abortSignal: AbortSignal,
+): () => void {
     const onAbort = () => {
-        void session.cancel().catch(() => undefined);
+        void server.cancelSession({ sessionId }).catch(() => undefined);
     };
     if (abortSignal.aborted) {
         onAbort();
@@ -49,63 +53,50 @@ function bindAbort(session: AgentSession, abortSignal: AbortSignal): () => void
 }
 
 export async function* streamTurnContent(
-    session: AgentSession,
+    server: AgentChatServer,
+    sessionId: string,
     foldState: PeerThreadFoldState,
     options: StreamTurnOptions,
     abortSignal: AbortSignal,
     groupRootBaseline?: readonly string[],
     /**
-     * Called once with the gateway-assigned turn ID as soon as it becomes
-     * available (after the first `turn.created` SSE event). Use this to
-     * reconcile the locally-generated optimistic ID with the real gateway ID
-     * so that edit/retry can find the turn in `buildSnapshotBeforeTurn`.
+     * Called once with the turn ID as soon as it becomes available (first
+     * `turn.created` SSE event). Use this to reconcile the locally-generated
+     * optimistic ID with the real turn ID.
      */
     onTurnIdAvailable?: (turnId: string) => void,
 ): AsyncGenerator<TurnStreamUpdate> {
-    const turn = session.prepareTurn({
-        input: buildTurnInput(options),
-        previousTurnId: options.previousTurnId ?? "auto",
-    });
-
-    const onAbort = bindAbort(session, abortSignal);
+    const onAbort = bindAbort(server, sessionId, abortSignal);
     if (abortSignal.aborted) {
         return;
     }
 
     let turnIdNotified = false;
-    const notifyTurnIdIfAvailable = () => {
-        if (!turnIdNotified && turn.id != null) {
-            onTurnIdAvailable?.(turn.id);
+    const notifyTurnId = (turnId: string) => {
+        if (!turnIdNotified) {
+            onTurnIdAvailable?.(turnId);
             turnIdNotified = true;
         }
     };
 
+    const stream: AsyncIterable<TurnStreamData> = server.prepareAndExecuteTurn({
+        sessionId,
+        input: buildTurnInput(options),
+        previousTurnId: options.previousTurnId ?? "auto",
+        abortSignal,
+        ...(options.headers != null ? { headers: options.headers } : {}),
+    });
+
     try {
         for await (const update of streamTurnEvents(
-            turn.execute(
-                { stream: true },
-                {
-                    abortSignal,
-                    ...(options.headers != null ? { headers: options.headers } : {}),
-                },
-            ),
+            stream,
             foldState,
             groupRootBaseline,
+            notifyTurnId,
         )) {
-            // After the first `turn.created` event, `turn.id` is set.
-            // Notify BEFORE yielding so the caller can update its tracking
-            // before the snapshot is written with the stream update.
-            notifyTurnIdIfAvailable();
             yield update;
         }
-        // Handle streams that complete without yielding any content.
-        notifyTurnIdIfAvailable();
     } catch (error) {
-        // Error streams often throw on `turn.done` (status=error) without ever
-        // yielding content (e.g. model not servable). `turn.id` is still set
-        // after `turn.created` — notify so edit/retry can resolve the turn.
-        // Same for AbortError after create: keep local ids aligned with gateway.
-        notifyTurnIdIfAvailable();
         if (error instanceof Error && error.name === "AbortError") {
             return;
         }
@@ -117,14 +108,22 @@ export async function* streamTurnContent(
 
 /** TODO: wire `afterSequenceNumber` from the last ingested stream event to skip replay on reconnect. */
 export async function* resumeTurnStream(
-    turn: Turn,
+    server: AgentChatServer,
+    sessionId: string,
+    turnId: string,
     foldState: PeerThreadFoldState,
     abortSignal: AbortSignal,
     afterSequenceNumber?: number,
     groupRootBaseline?: readonly string[],
 ): AsyncGenerator<TurnStreamUpdate> {
+    if (server.subscribeToTurn == null) {
+        throw new Error(
+            "resumeTurnStream requires AgentChatServer.subscribeToTurn",
+        );
+    }
+
     const onAbort = () => {
-        void turn.session.cancel().catch(() => undefined);
+        void server.cancelSession({ sessionId }).catch(() => undefined);
     };
     if (abortSignal.aborted) {
         onAbort();
@@ -134,10 +133,14 @@ export async function* resumeTurnStream(
 
     try {
         yield* streamTurnEvents(
-            turn.stream(
-                afterSequenceNumber != null ? { afterSequenceNumber } : {},
-                { abortSignal },
-            ),
+            server.subscribeToTurn({
+                sessionId,
+                turnId,
+                ...(afterSequenceNumber != null
+                    ? { afterSequenceNumber }
+                    : {}),
+                abortSignal,
+            }),
             foldState,
             groupRootBaseline,
         );

package/src/toolApproval.ts

@@ -8,7 +8,7 @@ import type {
     ToolApprovalRequiredEvent,
     Turn,
     UserToolApprovalEvent,
-} from "truefoundry-gateway-sdk/agents";
+} from "./server/index.js";
 
 import { ROOT_THREAD_ID } from "./constants.js";
 import type { ToolApprovalMessageCustomMetadata } from "./messageCustomMetadata.js";
@@ -148,10 +148,10 @@ export function findApprovalRequiredInTurn(
     if (turn.state.status !== "done") {
         return undefined;
     }
-    return turn.state.requiredActions?.find(
-        (action): action is ToolApprovalRequiredEvent =>
-            action.type === "tool.approval_required",
+    const found = turn.state.requiredActions?.find(
+        (action) => action.type === "tool.approval_required",
     );
+    return found?.type === "tool.approval_required" ? found : undefined;
 }
 
 function toolCallPartHasPendingApproval(part: ToolCallPart): boolean {

package/src/toolResponse.ts

@@ -9,7 +9,7 @@ import type {
     Turn,
     TurnInputItem,
     UserToolResponseEvent,
-} from "truefoundry-gateway-sdk/agents";
+} from "./server/index.js";
 
 import { ROOT_THREAD_ID } from "./constants.js";
 import {
@@ -85,10 +85,10 @@ export function findResponseRequiredInTurn(
     if (turn.state.status !== "done") {
         return undefined;
     }
-    return turn.state.requiredActions?.find(
-        (action): action is ToolResponseRequiredEvent =>
-            action.type === "tool.response_required",
+    const found = turn.state.requiredActions?.find(
+        (action) => action.type === "tool.response_required",
     );
+    return found?.type === "tool.response_required" ? found : undefined;
 }
 
 function applyToolResponseToToolCall(

package/src/truefoundryExtras.ts

@@ -1,5 +1,5 @@
 import { createRuntimeExtras } from "@assistant-ui/core/internal";
-import type { McpAuthRequiredEvent } from "truefoundry-gateway-sdk/agents";
+import type { McpAuthRequiredEvent } from "./server/index.js";
 
 import type { AgentSpec, AgentSpecUpdate } from "./private/agentSpec.js";
 import type { PendingApproval, PendingToolResponse } from "./collectPending.js";

package/src/truefoundryOwnedSessionsThreadListAdapter.test.ts

@@ -1,53 +1,54 @@
 import { describe, expect, it, vi } from "vitest";
 
-import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
+import type { AgentChatServer, Session } from "./server/index.js";
 
 import { createTrueFoundryOwnedSessionsThreadListAdapter } from "./truefoundryOwnedSessionsThreadListAdapter.js";
 
-function mockNamedSession(id: string, title: string, updatedAt: string) {
+function mockNamedSession(id: string, title: string, updatedAt: string): Session {
     return {
-        type: "session" as const,
         id,
         agentName: "my-agent",
         title,
-        createdBySubject: { type: "user" as const, id: "u1" },
         createdAt: updatedAt,
         updatedAt,
+        isMutable: false,
     };
 }
 
-function mockDraftSession(id: string, title: string | undefined, updatedAt: string) {
+function mockDraftSession(
+    id: string,
+    title: string | undefined,
+    updatedAt: string,
+): Session {
     return {
-        type: "session/draft" as const,
         id,
         agentSpec: { model: { name: "anthropic/claude-sonnet-4-6" } },
         title,
-        createdBySubject: { type: "user" as const, id: "u1" },
         createdAt: updatedAt,
         updatedAt,
+        isMutable: true,
     };
 }
 
+function mockServer(partial: Partial<AgentChatServer>): AgentChatServer {
+    return partial as AgentChatServer;
+}
+
 describe("createTrueFoundryOwnedSessionsThreadListAdapter", () => {
     it("lists owned sessions (named + draft) with pagination cursor", async () => {
-        const listOwnedSessions = vi.fn().mockResolvedValue({
+        const listSessions = vi.fn().mockResolvedValue({
             data: [
                 mockNamedSession("s1", "Named chat", "2026-06-30T12:00:00.000Z"),
                 mockDraftSession("d1", "Draft chat", "2026-06-30T11:00:00.000Z"),
             ],
-            response: {
-                pagination: { nextPageToken: "page-2", limit: 20 },
-            },
+            nextPageToken: "page-2",
         });
-        const privateClient = {
-            listOwnedSessions,
-            getDraftSession: vi.fn(),
-        } as unknown as PrivateAgentSessionClient;
+        const server = mockServer({ listSessions, getSession: vi.fn() });
 
-        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ privateClient });
+        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ server });
         const result = await adapter.list();
 
-        expect(listOwnedSessions).toHaveBeenCalledWith(
+        expect(listSessions).toHaveBeenCalledWith(
             expect.objectContaining({
                 limit: 20,
                 pageToken: undefined,
@@ -72,28 +73,24 @@ describe("createTrueFoundryOwnedSessionsThreadListAdapter", () => {
     });
 
     it("falls back to model name for untitled drafts", async () => {
-        const listOwnedSessions = vi.fn().mockResolvedValue({
+        const listSessions = vi.fn().mockResolvedValue({
             data: [mockDraftSession("d1", undefined, "2026-06-30T11:00:00.000Z")],
-            response: { pagination: { limit: 20 } },
         });
-        const privateClient = {
-            listOwnedSessions,
-            getDraftSession: vi.fn(),
-        } as unknown as PrivateAgentSessionClient;
+        const server = mockServer({ listSessions, getSession: vi.fn() });
 
-        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ privateClient });
+        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ server });
         const result = await adapter.list();
 
         expect(result.threads[0]?.title).toBe("anthropic/claude-sonnet-4-6");
     });
 
     it("throws on initialize because the adapter is read-only", async () => {
-        const privateClient = {
-            listOwnedSessions: vi.fn(),
-            getDraftSession: vi.fn(),
-        } as unknown as PrivateAgentSessionClient;
+        const server = mockServer({
+            listSessions: vi.fn(),
+            getSession: vi.fn(),
+        });
 
-        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ privateClient });
+        const adapter = createTrueFoundryOwnedSessionsThreadListAdapter({ server });
 
         await expect(adapter.initialize("local")).rejects.toThrow(/read-only/);
     });

package/src/truefoundryOwnedSessionsThreadListAdapter.ts

@@ -1,34 +1,33 @@
 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 type { AgentChatServer, Session } from "./server/types.js";
 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);
+function ownedSessionTitle(session: Session): string {
+    if (session.isMutable && session.agentSpec != null) {
+        return draftSessionTitle({
+            title: session.title,
+            agentSpec: session.agentSpec,
+        });
     }
-    return session.title ?? session.agentName;
+    return session.title ?? session.agentName ?? session.id;
 }
 
 /**
- * Read-only thread-list adapter backed by `PrivateAgentSessionClient.listOwnedSessions`.
- * Returns every session the caller owns (named + draft), newest first.
+ * Read-only thread-list adapter backed by `AgentChatServer.listSessions`.
+ * Hosts that previously used listOwnedSessions should filter in their server impl.
  */
 export function createTrueFoundryOwnedSessionsThreadListAdapter(options: {
-    privateClient: PrivateAgentSessionClient;
+    server: AgentChatServer;
 }): RemoteThreadListAdapter {
-    const { privateClient } = options;
+    const { server } = options;
 
     return {
         async list({ after } = {}) {
-            const page = await privateClient.listOwnedSessions({
+            const page = await server.listSessions({
                 limit: THREAD_LIST_PAGE_SIZE,
                 pageToken: after,
                 startTimestamp: sessionListStartTimestamp(),
@@ -41,7 +40,7 @@ export function createTrueFoundryOwnedSessionsThreadListAdapter(options: {
             }));
             return {
                 threads,
-                nextCursor: page.response.pagination.nextPageToken ?? undefined,
+                nextCursor: page.nextPageToken ?? undefined,
             };
         },
 
@@ -52,16 +51,12 @@ export function createTrueFoundryOwnedSessionsThreadListAdapter(options: {
         },
 
         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,
-            });
+            const session = await server.getSession({ sessionId: remoteId });
             return {
                 status: "regular" as const,
-                remoteId: draft.id,
-                title: draftSessionTitle(draft),
-                lastMessageAt: new Date(draft.updatedAt),
+                remoteId: session.id,
+                title: ownedSessionTitle(session),
+                lastMessageAt: new Date(session.updatedAt),
             };
         },
 

package/src/truefoundryThreadListAdapter.test.ts

@@ -1,32 +1,30 @@
 import { describe, expect, it, vi } from "vitest";
 
-import type { AgentSessionClient } from "truefoundry-gateway-sdk/agents";
+import type { AgentChatServer, Session } from "./server/index.js";
 
 import { createTrueFoundryThreadListAdapter } from "./truefoundryThreadListAdapter.js";
 
-function mockSession(id: string, title: string, updatedAt: string) {
+function mockSession(id: string, title: string, updatedAt: string): Session {
     return {
         id,
         title,
         updatedAt,
+        createdAt: updatedAt,
+        isMutable: false,
     };
 }
 
-function mockListSessionsPage(
-    sessions: ReturnType<typeof mockSession>[],
-    nextPageToken?: string,
-) {
+function mockListSessionsPage(sessions: Session[], nextPageToken?: string) {
     return {
         data: sessions,
-        response: {
-            pagination: {
-                nextPageToken,
-                limit: 20,
-            },
-        },
+        ...(nextPageToken != null ? { nextPageToken } : {}),
     };
 }
 
+function mockServer(partial: Partial<AgentChatServer>): AgentChatServer {
+    return partial as AgentChatServer;
+}
+
 describe("createTrueFoundryThreadListAdapter", () => {
     it("lists the first page with limit and returns nextCursor", async () => {
         const listSessions = vi.fn().mockResolvedValue(
@@ -35,9 +33,9 @@ describe("createTrueFoundryThreadListAdapter", () => {
                 "page-2",
             ),
         );
-        const client = { listSessions } as unknown as AgentSessionClient;
+        const server = mockServer({ listSessions });
         const adapter = createTrueFoundryThreadListAdapter({
-            client,
+            server,
             agentName: "my-agent",
         });
 
@@ -68,9 +66,9 @@ describe("createTrueFoundryThreadListAdapter", () => {
                 [mockSession("s2", "Second", "2026-06-29T10:00:00.000Z")],
             ),
         );
-        const client = { listSessions } as unknown as AgentSessionClient;
+        const server = mockServer({ listSessions });
         const adapter = createTrueFoundryThreadListAdapter({
-            client,
+            server,
             agentName: "my-agent",
         });
 
@@ -90,9 +88,9 @@ describe("createTrueFoundryThreadListAdapter", () => {
         const listSessions = vi.fn().mockResolvedValue(
             mockListSessionsPage([mockSession("s1", "Only", "2026-06-30T10:00:00.000Z")]),
         );
-        const client = { listSessions } as unknown as AgentSessionClient;
+        const server = mockServer({ listSessions });
         const adapter = createTrueFoundryThreadListAdapter({
-            client,
+            server,
             agentName: "my-agent",
         });
 

package/src/truefoundryThreadListAdapter.ts

@@ -1,20 +1,20 @@
 import type { RemoteThreadListAdapter } from "@assistant-ui/core";
-import type { AgentSessionClient } from "truefoundry-gateway-sdk/agents";
 
+import type { AgentChatServer } from "./server/types.js";
 import { getSession } from "./sessions.js";
 import { sessionListStartTimestamp } from "./sessionListStartTimestamp.js";
 
 const THREAD_LIST_PAGE_SIZE = 20;
 
 export function createTrueFoundryThreadListAdapter(options: {
-    client: AgentSessionClient;
+    server: AgentChatServer;
     agentName: string;
 }): RemoteThreadListAdapter {
-    const { client, agentName } = options;
+    const { server, agentName } = options;
 
     return {
         async list({ after } = {}) {
-            const page = await client.listSessions({
+            const page = await server.listSessions({
                 agentName,
                 limit: THREAD_LIST_PAGE_SIZE,
                 pageToken: after,
@@ -28,17 +28,17 @@ export function createTrueFoundryThreadListAdapter(options: {
             }));
             return {
                 threads,
-                nextCursor: page.response.pagination.nextPageToken ?? undefined,
+                nextCursor: page.nextPageToken ?? undefined,
             };
         },
 
         async initialize(_threadId: string) {
-            const session = await client.createSession({ agentName });
+            const session = await server.createSession({ agentName });
             return { remoteId: session.id, externalId: undefined };
         },
 
         async fetch(remoteId) {
-            const session = await getSession(client, remoteId);
+            const session = await getSession(server, remoteId);
             return {
                 status: "regular" as const,
                 remoteId: session.id,

package/src/turnEventHelpers.ts

@@ -1,4 +1,4 @@
-import type { Turn } from "truefoundry-gateway-sdk/agents";
+import type { Turn } from "./server/index.js";
 
 import {
     buildMcpAuthTextParts,

package/src/types.ts

@@ -6,10 +6,9 @@ import type {
     RealtimeVoiceAdapter,
     SpeechSynthesisAdapter,
 } from "@assistant-ui/core";
-import type { AgentSessionClient } from "truefoundry-gateway-sdk/agents";
-import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
 
 import type { AgentSpec } from "./private/agentSpec.js";
+import type { AgentChatServer } from "./server/types.js";
 
 export type NamedAgentConfig = {
     mode: "named";
@@ -25,7 +24,7 @@ export type DraftAgentConfig = {
 export type TrueFoundryAgentConfig = NamedAgentConfig | DraftAgentConfig;
 
 type TrueFoundryAgentRuntimeBaseOptions = ExternalStoreSharedOptions & {
-    client: AgentSessionClient;
+    server: AgentChatServer;
     initialSessionId?: string | undefined;
     threadId?: string | undefined;
     onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
@@ -47,16 +46,10 @@ export type UseTrueFoundryAgentRuntimeOptions = TrueFoundryAgentRuntimeBaseOptio
     agent?: TrueFoundryAgentConfig | undefined;
     /** Legacy named-agent shorthand. Prefer `agent: { mode: "named", agentName }`. */
     agentName?: string | undefined;
-    /**
-     * Required when `agent.mode === "draft"`. Also used for sandbox file downloads
-     * (named or draft).
-     */
-    privateClient?: PrivateAgentSessionClient | undefined;
 };
 
 export type ResolvedTrueFoundryAgentRuntimeOptions = TrueFoundryAgentRuntimeBaseOptions & {
     agent: TrueFoundryAgentConfig;
-    privateClient?: PrivateAgentSessionClient | undefined;
 };
 
 export function resolveTrueFoundryAgentConfig(
@@ -81,15 +74,8 @@ export function resolveTrueFoundryAgentRuntimeOptions(
 ): ResolvedTrueFoundryAgentRuntimeOptions {
     const agent = resolveTrueFoundryAgentConfig(options);
 
-    if (agent.mode === "draft" && options.privateClient == null) {
-        throw new Error(
-            "Draft agent mode requires a `privateClient` PrivateAgentSessionClient.",
-        );
-    }
-
     return {
         ...options,
         agent,
-        privateClient: options.privateClient,
     };
 }