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

package/src/loadSessionSnapshot.test.ts

@@ -1,4 +1,4 @@
-import type { AgentSessionClient } from "truefoundry-gateway-sdk/agents";
+import type { AgentChatServer } from "./server/index.js";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
 import { createEmptySessionSnapshot } from "./sessionSnapshot.js";
@@ -19,7 +19,7 @@ import { buildSnapshotFromSessionEvents } from "./convertTurnMessages.js";
 import { loadSessionSnapshot } from "./loadSessionSnapshot.js";
 import { getSession } from "./sessions.js";
 
-const mockClient = {} as AgentSessionClient;
+const mockServer = {} as AgentChatServer;
 
 describe("loadSessionSnapshot", () => {
     beforeEach(() => {
@@ -33,8 +33,8 @@ describe("loadSessionSnapshot", () => {
         );
 
         const [first, second] = await Promise.all([
-            loadSessionSnapshot(mockClient, "session-1"),
-            loadSessionSnapshot(mockClient, "session-1"),
+            loadSessionSnapshot(mockServer, "session-1"),
+            loadSessionSnapshot(mockServer, "session-1"),
         ]);
 
         expect(first).toBe(second);
@@ -48,8 +48,8 @@ describe("loadSessionSnapshot", () => {
             createEmptySessionSnapshot(),
         );
 
-        await loadSessionSnapshot(mockClient, "session-1");
-        await loadSessionSnapshot(mockClient, "session-1");
+        await loadSessionSnapshot(mockServer, "session-1");
+        await loadSessionSnapshot(mockServer, "session-1");
 
         expect(getSession).toHaveBeenCalledTimes(2);
         expect(buildSnapshotFromSessionEvents).toHaveBeenCalledTimes(2);
@@ -62,10 +62,11 @@ describe("loadSessionSnapshot", () => {
         );
 
         const onProgress = vi.fn();
-        await loadSessionSnapshot(mockClient, "session-1", undefined, onProgress);
+        await loadSessionSnapshot(mockServer, "session-1", onProgress);
 
         expect(buildSnapshotFromSessionEvents).toHaveBeenCalledWith(
-            {},
+            mockServer,
+            "session-1",
             onProgress,
         );
     });

package/src/loadSessionSnapshot.ts

@@ -1,7 +1,7 @@
-import type { AgentSessionClient } from "truefoundry-gateway-sdk/agents";
+import type { AgentChatServer } from "./server/types.js";
 
 import { buildSnapshotFromSessionEvents } from "./convertTurnMessages.js";
-import { getSession, type GetSessionOptions } from "./sessions.js";
+import { getSession } from "./sessions.js";
 import type { SessionSnapshot } from "./sessionSnapshot.js";
 
 const inflightBySessionId = new Map<string, Promise<SessionSnapshot>>();
@@ -15,25 +15,20 @@ const inflightBySessionId = new Map<string, Promise<SessionSnapshot>>();
  * can update the UI progressively while history is being processed.
  */
 export function loadSessionSnapshot(
-    client: AgentSessionClient,
+    server: AgentChatServer,
     sessionId: string,
-    sessionOptions?: GetSessionOptions,
     onProgress?: (snap: SessionSnapshot) => void,
 ): Promise<SessionSnapshot> {
-    const cacheKey =
-        sessionOptions?.privateClient != null ? `draft:${sessionId}` : sessionId;
-    let inflight = inflightBySessionId.get(cacheKey);
+    let inflight = inflightBySessionId.get(sessionId);
     if (inflight == null) {
-        inflight = getSession(client, sessionId, sessionOptions)
-            .then((session) =>
-                buildSnapshotFromSessionEvents(session, onProgress),
-            )
+        inflight = getSession(server, sessionId)
+            .then(() => buildSnapshotFromSessionEvents(server, sessionId, onProgress))
             .finally(() => {
-                if (inflightBySessionId.get(cacheKey) === inflight) {
-                    inflightBySessionId.delete(cacheKey);
+                if (inflightBySessionId.get(sessionId) === inflight) {
+                    inflightBySessionId.delete(sessionId);
                 }
             });
-        inflightBySessionId.set(cacheKey, inflight);
+        inflightBySessionId.set(sessionId, inflight);
     }
     return inflight;
 }

package/src/mcpAuth.ts

@@ -1,5 +1,5 @@
 import type { MessageStatus } from "@assistant-ui/core";
-import type { McpAuthRequiredEvent, TurnStateDone } from "truefoundry-gateway-sdk/agents";
+import type { McpAuthRequiredEvent, TurnStateDone } from "./server/index.js";
 
 import type { AssistantContentPart } from "./modelMessageContent.js";
 import type { McpAuthMessageCustomMetadata } from "./messageCustomMetadata.js";
@@ -28,9 +28,12 @@ export function buildMcpAuthTextParts(
 export function findMcpAuthRequired(
     requiredActions: TurnStateDone["requiredActions"] | undefined,
 ): McpAuthRequiredEvent | undefined {
-    return requiredActions?.find(
-        (action): action is McpAuthRequiredEvent => action.type === "mcp.auth_required",
+    const found = requiredActions?.find(
+        (action) => action.type === "mcp.auth_required",
     );
+    return found?.type === "mcp.auth_required"
+        ? (found as McpAuthRequiredEvent)
+        : undefined;
 }
 
 export function mcpAuthAssistantStatus(): MessageStatus {

package/src/messageCustomMetadata.ts

@@ -1,4 +1,4 @@
-import type { McpAuthRequiredEvent } from "truefoundry-gateway-sdk/agents";
+import type { McpAuthRequiredEvent } from "./server/index.js";
 
 import type { SubAgentCustomMetadata } from "./foldPeerThreads.js";
 import { TOOL_APPROVAL_THREAD_ID_CUSTOM_KEY } from "./toolApproval.js";

package/src/modelMessageContent.ts

@@ -1,5 +1,5 @@
 import type { ThreadAssistantMessagePart } from "@assistant-ui/core";
-import type { ModelMessageEvent } from "truefoundry-gateway-sdk/agents";
+import type { ModelMessageEvent } from "./server/index.js";
 import type { PendingResponseRef } from "./foldPeerThreads.js";
 import { extractImagePartsFromModelMessage } from "./modelMessageImageContent.js";
 

package/src/modelMessageImageContent.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "vitest";
-import type { ModelMessageEvent } from "truefoundry-gateway-sdk/agents";
+import type { ModelMessageEvent } from "./server/index.js";
 
 import { buildAssistantContent } from "./modelMessageContent.js";
 import { mergeStreamEventDelta } from "./modelMessageImageContent.js";

package/src/modelMessageImageContent.ts

@@ -1,10 +1,11 @@
 import type { CompleteAttachment } from "@assistant-ui/core";
-import type {
-    ModelMessageEvent,
-    TurnEvent,
-    TurnStreamingEvent,
-} from "truefoundry-gateway-sdk/agents";
-import { isEventDelta, mergeEventDelta } from "truefoundry-gateway-sdk/agents";
+import {
+    isEventDelta,
+    mergeEventDelta,
+    type ModelMessageEvent,
+    type TurnEvent,
+    type TurnStreamingEvent,
+} from "./server/index.js";
 
 import type { AssistantContentPart } from "./modelMessageContent.js";
 

package/src/plugins/truefoundry-agent-server-adapter/index.ts

@@ -0,0 +1,285 @@
+import { AgentSessionClient } from "truefoundry-gateway-sdk/agents";
+import type { AgentSession } from "truefoundry-gateway-sdk/agents";
+import { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
+import type { AgentDraftSession } from "truefoundry-gateway-sdk/agents/private";
+import type {
+    AgentChatServer,
+    AgentSpec,
+    ListResult,
+    Session,
+    Turn,
+    TurnInputItem,
+    PreviousTurnIdInput,
+} from "../../server/types.js";
+import type {
+    TurnEvent,
+    TurnStreamData,
+    SessionEventItem,
+} from "../../server/events.js";
+
+type GwSession = AgentSession | AgentDraftSession;
+
+export type CreateTrueFoundryChatServerOptions = {
+    apiKey: string;
+    baseUrl: string;
+    /** Optional override — otherwise constructed from apiKey/baseUrl. */
+    client?: AgentSessionClient;
+    privateClient?: PrivateAgentSessionClient;
+    deleteSession?: (req: { sessionId: string }) => Promise<void>;
+};
+
+export type TrueFoundryChatServer = AgentChatServer & {
+    /** Escape hatch for hosts that still need raw gateway clients. */
+    getGatewayClients(): {
+        client: AgentSessionClient;
+        privateClient: PrivateAgentSessionClient;
+    };
+};
+
+function isDraft(session: GwSession): session is AgentDraftSession {
+    return (session as AgentDraftSession).type === "session/draft";
+}
+
+function toSession(raw: GwSession): Session {
+    const mutable = isDraft(raw);
+    return {
+        id: raw.id,
+        title: raw.title,
+        agentName: raw.agentName,
+        ...(mutable ? { agentSpec: raw.agentSpec as AgentSpec } : {}),
+        isMutable: mutable,
+        createdAt: raw.createdAt,
+        updatedAt: raw.updatedAt,
+    };
+}
+
+function toTurn(raw: {
+    id: string;
+    sessionId: string;
+    previousTurnId?: string | null;
+    input?: TurnInputItem[];
+    state: Turn["state"];
+    createdAt: string;
+}): Turn {
+    return {
+        id: raw.id,
+        sessionId: raw.sessionId,
+        previousTurnId: raw.previousTurnId,
+        input: raw.input as TurnInputItem[] | undefined,
+        state: raw.state as Turn["state"],
+        createdAt: raw.createdAt,
+    };
+}
+
+async function toListResult<TIn, TOut>(
+    page: {
+        data: TIn[];
+        response?: { pagination?: { nextPageToken?: string } };
+        hasNextPage?: () => boolean;
+    },
+    map: (item: TIn) => TOut,
+): Promise<ListResult<TOut>> {
+    const nextPageToken = page.response?.pagination?.nextPageToken;
+    return {
+        data: page.data.map(map),
+        ...(nextPageToken != null && nextPageToken !== ""
+            ? { nextPageToken }
+            : {}),
+    };
+}
+
+/**
+ * Wraps TrueFoundry gateway clients into a flat `AgentChatServer`.
+ * Named vs draft routing is fully internal — an in-memory session-type cache
+ * (populated by createSession/listSessions) determines which gateway client
+ * to call. No try/catch fallback, no double network calls.
+ */
+export function createTrueFoundryChatServer(
+    opts: CreateTrueFoundryChatServerOptions,
+): TrueFoundryChatServer {
+    const gatewayOpts = { apiKey: opts.apiKey, baseUrl: opts.baseUrl };
+    const client = opts.client ?? new AgentSessionClient(gatewayOpts);
+    const privateClient =
+        opts.privateClient ?? new PrivateAgentSessionClient(gatewayOpts);
+
+    const sessionTypeCache = new Map<string, boolean>();
+
+    function cacheSessionType(session: Session): void {
+        sessionTypeCache.set(session.id, session.isMutable);
+    }
+
+    function getSessionObj(sessionId: string): Promise<GwSession> {
+        const isMutable = sessionTypeCache.get(sessionId);
+        if (isMutable === true) {
+            return privateClient.getDraftSession({ draftSessionId: sessionId });
+        }
+        if (isMutable === false) {
+            return client.getSession({ sessionId });
+        }
+        throw new Error(
+            `Cannot resolve session "${sessionId}": session type not cached. ` +
+                `Ensure createSession or listSessions was called first.`,
+        );
+    }
+
+    const server: TrueFoundryChatServer = {
+        async createSession(req) {
+            if (req.agentSpec != null) {
+                const draft = await privateClient.createDraftSession({
+                    agentSpec: req.agentSpec as never,
+                    ...(req.agentName != null ? { agentName: req.agentName } : {}),
+                });
+                const session = toSession(draft);
+                cacheSessionType(session);
+                return session;
+            }
+            if (req.agentName != null) {
+                const named = await client.createSession({
+                    agentName: req.agentName,
+                });
+                const session = toSession(named);
+                cacheSessionType(session);
+                return session;
+            }
+            throw new Error("createSession requires agentName and/or agentSpec");
+        },
+
+        async listSessions(req) {
+            const page = await privateClient.listOwnedSessions({
+                limit: req?.limit,
+                order: req?.order,
+                pageToken: req?.pageToken,
+                startTimestamp: req?.startTimestamp,
+                ...(req?.agentName != null ? { agentName: req.agentName } : {}),
+            });
+            const result = await toListResult(page, toSession);
+            for (const session of result.data) {
+                cacheSessionType(session);
+            }
+            return result;
+        },
+
+        async getSession({ sessionId }) {
+            const raw = await getSessionObj(sessionId);
+            const session = toSession(raw);
+            cacheSessionType(session);
+            return session;
+        },
+
+        async updateSession(req) {
+            const raw = await getSessionObj(req.sessionId);
+            if (!isDraft(raw)) {
+                throw new Error(
+                    "updateSession: session is not mutable (isMutable=false)",
+                );
+            }
+            if (req.agentSpec != null) {
+                await raw.update({ agentSpec: req.agentSpec as never });
+            }
+            return toSession(raw);
+        },
+
+        prepareAndExecuteTurn(req: {
+            sessionId: string;
+            input?: TurnInputItem[];
+            previousTurnId?: PreviousTurnIdInput;
+            abortSignal?: AbortSignal;
+            headers?: Record<string, string>;
+        }): AsyncIterable<TurnStreamData> {
+            return (async function* () {
+                const session = await getSessionObj(req.sessionId);
+                const prepared = session.prepareTurn({
+                    input: req.input,
+                    previousTurnId: req.previousTurnId ?? "auto",
+                });
+                yield* prepared.execute(
+                    { stream: true },
+                    {
+                        ...(req.abortSignal != null
+                            ? { abortSignal: req.abortSignal }
+                            : {}),
+                        ...(req.headers != null ? { headers: req.headers } : {}),
+                    },
+                ) as AsyncIterable<TurnStreamData>;
+            })();
+        },
+
+        async cancelSession({ sessionId }) {
+            await (await getSessionObj(sessionId)).cancel();
+        },
+
+        async deleteSession({ sessionId }) {
+            if (opts.deleteSession == null) {
+                throw new Error(
+                    "deleteSession is not on the gateway SDK. Pass deleteSession to createTrueFoundryChatServer.",
+                );
+            }
+            await opts.deleteSession({ sessionId });
+        },
+
+        async listTurns({ sessionId, limit, pageToken, order }) {
+            const raw = await getSessionObj(sessionId);
+            const page = await raw.listTurns({
+                ...(limit != null ? { limit } : {}),
+                ...(pageToken != null ? { pageToken } : {}),
+                ...(order != null ? { order } : {}),
+            });
+            return toListResult(page, (turn) => toTurn(turn));
+        },
+
+        async getTurn({ sessionId, turnId }) {
+            const raw = await getSessionObj(sessionId);
+            return toTurn(await raw.getTurn({ turnId }));
+        },
+
+        async listEvents({ sessionId, pageToken, lastTurnId, limit }) {
+            const raw = await getSessionObj(sessionId);
+            const page = await raw.listEvents({
+                ...(limit != null ? { limit } : {}),
+                ...(pageToken != null ? { pageToken } : {}),
+                ...(lastTurnId != null ? { lastTurnId } : {}),
+            });
+            return toListResult(
+                page,
+                (item) => item as SessionEventItem,
+            );
+        },
+
+        async listTurnEvents({ sessionId, turnId, limit, pageToken, order }) {
+            const raw = await getSessionObj(sessionId);
+            const turn = await raw.getTurn({ turnId });
+            const page = await turn.listEvents({
+                ...(limit != null ? { limit } : {}),
+                ...(pageToken != null ? { pageToken } : {}),
+                ...(order != null ? { order } : {}),
+            });
+            return toListResult(page, (event) => event as TurnEvent);
+        },
+
+        async *subscribeToTurn({
+            sessionId,
+            turnId,
+            afterSequenceNumber,
+            abortSignal,
+        }) {
+            const raw = await getSessionObj(sessionId);
+            const turn = await raw.getTurn({ turnId });
+            yield* turn.stream(
+                afterSequenceNumber != null ? { afterSequenceNumber } : {},
+                abortSignal != null ? { abortSignal } : {},
+            ) as AsyncIterable<TurnStreamData>;
+        },
+
+        async downloadSandboxFile(sandboxId, req) {
+            const response = await privateClient.downloadSandboxFile(
+                sandboxId,
+                req,
+            );
+            return await response.blob();
+        },
+
+        getGatewayClients: () => ({ client, privateClient }),
+    };
+
+    return server;
+}

package/src/private/agentSpec.ts

@@ -1,7 +1,12 @@
-import type { TruefoundryGatewayApi } from "truefoundry-gateway-sdk";
+import type { AgentSpec } from "../server/types.js";
 
-export type AgentSpec = TruefoundryGatewayApi.AgentSpec;
-export type DraftSession = TruefoundryGatewayApi.DraftSession;
+export type { AgentSpec } from "../server/types.js";
+
+/** @deprecated Use Session with isMutable: true instead. Kept for title helper. */
+export type DraftSession = {
+    title?: string | null;
+    agentSpec: AgentSpec;
+};
 
 export type AgentSpecUpdate = {
     instructions?: string;

package/src/private/draftSessionBridge.ts

@@ -1,7 +1,5 @@
-import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
-
+import type { AgentChatServer } from "../server/types.js";
 import type { AgentSpec } from "./agentSpec.js";
-import { getGatewayFromPrivateClient } from "./getGatewayFromPrivateClient.js";
 
 export const DRAFT_SESSION_LAST_UPDATED_AT_HEADER = "x-tfy-session-last-updated-at";
 
@@ -11,22 +9,25 @@ export type DraftSessionBridge = {
 };
 
 export function createDraftSessionBridge(
-    privateClient: PrivateAgentSessionClient,
+    server: AgentChatServer,
 ): DraftSessionBridge {
     return {
         async getDraftAgentSpec(draftSessionId) {
-            const draft = await privateClient.getDraftSession({ draftSessionId });
-            return draft.agentSpec;
+            const session = await server.getSession({ sessionId: draftSessionId });
+            if (session.agentSpec == null) {
+                throw new Error(
+                    `Session ${draftSessionId} has no agentSpec (isMutable=${session.isMutable}).`,
+                );
+            }
+            return session.agentSpec;
         },
 
         async syncAgentSpec(draftSessionId, agentSpec) {
-            // PrivateAgentSessionClient does not wrap update yet — use the low-level client.
-            const gateway = getGatewayFromPrivateClient(privateClient);
-            const response = await gateway.agents.private.draftSessions.update(
-                draftSessionId,
-                { agentSpec },
-            );
-            return response.data.updatedAt;
+            const updated = await server.updateSession({
+                sessionId: draftSessionId,
+                agentSpec,
+            });
+            return updated.updatedAt;
         },
     };
 }