@truefoundry/trueforge-assistant-ui-runtime 0.0.0 → 0.2.0-rc.0

package/src/createSubAgent.ts

@@ -0,0 +1,11 @@
+import type { ToolCall } from './server/index.js';
+
+export function isCreateSubAgentToolCall(toolCall: Pick<ToolCall, 'toolInfo' | 'function'>): boolean {
+  // Gateway may attach trueforge-system toolInfo on persisted turns, but streamed
+  // model.message tool calls often only carry function.name. Without the fallback,
+  // foldPeerThreads never nests child threads under the spawning tool call.
+  if (toolCall.toolInfo?.type === 'trueforge-system' && toolCall.toolInfo.name === 'create_sub_agent') {
+    return true;
+  }
+  return toolCall.function.name === 'create_sub_agent';
+}

package/src/draft/agentSpec.ts

@@ -0,0 +1,34 @@
+import type { AgentSpec } from '../server/types.js';
+
+export interface DraftSession {
+  title?: string | null;
+  agentSpec: AgentSpec;
+}
+
+/** Partial update — host fields flow through when `TSpec` is widened. */
+export type AgentSpecUpdate<TSpec extends AgentSpec = AgentSpec> = {
+  [K in keyof TSpec]?: K extends 'model'
+    ? Omit<Partial<TSpec['model']>, 'params'> & {
+        params?: Partial<NonNullable<TSpec['model']['params']>>;
+      }
+    : TSpec[K];
+};
+
+export function mergeAgentSpec<TSpec extends AgentSpec>(base: TSpec, update: AgentSpecUpdate<TSpec>): TSpec {
+  const { model: modelUpdate, ...rest } = update;
+
+  const next: TSpec = Object.assign({}, base, rest);
+
+  if (modelUpdate != null) {
+    next.model = Object.assign({}, base.model, modelUpdate, {
+      name: modelUpdate.name ?? base.model.name,
+      params: modelUpdate.params != null ? { ...base.model.params, ...modelUpdate.params } : base.model.params,
+    });
+  }
+
+  return next;
+}
+
+export function draftSessionTitle(draft: DraftSession): string {
+  return draft.title ?? draft.agentSpec.model.name;
+}

package/src/draft/draftSessionBridge.ts

@@ -0,0 +1,28 @@
+import type { AgentChatServer, AgentSpec } from '../server/types.js';
+
+export const DRAFT_SESSION_LAST_UPDATED_AT_HEADER = 'x-tfy-session-last-updated-at';
+
+export interface DraftSessionBridge {
+  syncAgentSpec: (draftSessionId: string, agentSpec: AgentSpec) => Promise<string>;
+  getDraftAgentSpec: (draftSessionId: string) => Promise<AgentSpec>;
+}
+
+export function createDraftSessionBridge(server: AgentChatServer): DraftSessionBridge {
+  return {
+    async getDraftAgentSpec(draftSessionId) {
+      const session = await server.getSession({ sessionId: draftSessionId });
+      if (session.agentSpec == null) {
+        throw new Error(`Session ${draftSessionId} has no agentSpec (isMutable=${String(session.isMutable)}).`);
+      }
+      return session.agentSpec;
+    },
+
+    async syncAgentSpec(draftSessionId, agentSpec) {
+      const updated = await server.updateSession({
+        sessionId: draftSessionId,
+        agentSpec,
+      });
+      return updated.updatedAt;
+    },
+  };
+}

package/src/draft/trueforgeDraftThreadListAdapter.ts

@@ -0,0 +1,73 @@
+import type { RemoteThreadListAdapter } from '@assistant-ui/core';
+
+import type { AgentChatServer, AgentSpec } from '../server/types.js';
+import { sessionListStartTimestamp } from '../sessionListStartTimestamp.js';
+import { sessionDisplayTitle, sessionToThreadMetadata } from '../sessionThreadMetadata.js';
+
+const THREAD_LIST_PAGE_SIZE = 20;
+
+export function createTrueForgeDraftThreadListAdapter(options: {
+  server: AgentChatServer;
+  defaultAgentSpec: AgentSpec;
+  getAgentSpec?: () => AgentSpec;
+  /** When set, filters `listSessions` by this agent id. Omit for all chats. */
+  listSessionsAgentId?: string;
+  /** Restrict history to sessions created by the authenticated subject. */
+  listSessionsCreatedByMe?: boolean;
+}): RemoteThreadListAdapter {
+  const { server, defaultAgentSpec, getAgentSpec, listSessionsAgentId, listSessionsCreatedByMe = false } = options;
+
+  return {
+    async list({ after } = {}) {
+      const page = await server.listSessions({
+        ...(listSessionsAgentId != null ? { agentId: listSessionsAgentId } : {}),
+        createdByMe: listSessionsCreatedByMe,
+        limit: THREAD_LIST_PAGE_SIZE,
+        ...(after == null ? {} : { pageToken: after }),
+        startTimestamp: sessionListStartTimestamp(),
+      });
+      const threads = page.data.map(session =>
+        sessionToThreadMetadata(session, sessionDisplayTitle(session, defaultAgentSpec)),
+      );
+      return {
+        threads,
+        nextCursor: page.nextPageToken ?? undefined,
+      };
+    },
+
+    async initialize() {
+      const draft = await server.createSession({
+        agentSpec: getAgentSpec?.() ?? defaultAgentSpec,
+      });
+      return { remoteId: draft.id, externalId: undefined };
+    },
+
+    async fetch(remoteId) {
+      const draft = await server.getSession({ sessionId: remoteId });
+      return sessionToThreadMetadata(draft, sessionDisplayTitle(draft, defaultAgentSpec));
+    },
+
+    async rename(remoteId, newTitle) {
+      if (typeof server.renameSession !== 'function') {
+        return;
+      }
+      await server.renameSession({ sessionId: remoteId, title: newTitle });
+    },
+    archive() {
+      return Promise.resolve();
+    },
+    unarchive() {
+      return Promise.resolve();
+    },
+    async delete(remoteId) {
+      if (typeof server.deleteSession !== 'function') {
+        return;
+      }
+      await server.deleteSession({ sessionId: remoteId });
+    },
+
+    generateTitle() {
+      return Promise.resolve(new ReadableStream());
+    },
+  };
+}

package/src/draft/useDraftAgentSpec.ts

@@ -0,0 +1,289 @@
+'use client';
+
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
+
+import type { AgentSpec } from '../server/types.js';
+import { mergeAgentSpec, type AgentSpecUpdate } from './agentSpec.js';
+import type { DraftSessionBridge } from './draftSessionBridge.js';
+
+const SPEC_SYNC_DEBOUNCE_MS = 400;
+
+export interface UseDraftAgentSpecOptions {
+  draftSessionId: string | undefined;
+  draftBridge: DraftSessionBridge | null;
+  defaultAgentSpec: AgentSpec;
+  onAgentSpecChange?: ((spec: AgentSpec) => void) | undefined;
+  onError?: ((error: unknown) => void) | undefined;
+}
+
+export interface UseDraftAgentSpecResult {
+  agentSpec: AgentSpec | null;
+  draftSessionId: string | undefined;
+  isSpecLoading: boolean;
+  isSpecSyncing: boolean;
+  specError: unknown;
+  updateAgentSpec: (update: AgentSpecUpdate) => void;
+  flushAgentSpec: () => Promise<void>;
+  adoptAgentSpec: (request: { agentSpec: AgentSpec; updatedAt?: string }) => void;
+  takeTurnHeaderTimestamp: () => Promise<string | undefined>;
+}
+
+export function useDraftAgentSpec({
+  draftSessionId,
+  draftBridge,
+  defaultAgentSpec,
+  onAgentSpecChange,
+  onError,
+}: UseDraftAgentSpecOptions): UseDraftAgentSpecResult {
+  const enabled = draftBridge != null;
+  const [agentSpec, setAgentSpec] = useState<AgentSpec>(defaultAgentSpec);
+  const [isSpecLoading, setIsSpecLoading] = useState(false);
+  const [isSpecSyncing, setIsSpecSyncing] = useState(false);
+  const [specError, setSpecError] = useState<unknown>(null);
+
+  const agentSpecRef = useRef(agentSpec);
+  agentSpecRef.current = agentSpec;
+
+  const syncTimeoutRef = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
+  const syncGenerationRef = useRef(0);
+  const loadedDraftIdRef = useRef<string | undefined>(undefined);
+  const localDirtyRef = useRef(false);
+  const lastUpdatedAtRef = useRef<string | undefined>(undefined);
+  const pendingFlushRef = useRef<(() => Promise<void>) | undefined>(undefined);
+  const inFlightFlushRef = useRef<Promise<void> | undefined>(undefined);
+  const activeDraftIdRef = useRef(draftSessionId);
+
+  // Invalidate sync state from the previous draft so a stale pending flush or
+  // stored updatedAt can't update the wrong draft or leak into the next
+  // turn's header. Bumping the generation makes any in-flight sync a no-op.
+  useEffect(() => {
+    const previousDraftId = activeDraftIdRef.current;
+    if (previousDraftId === draftSessionId) {
+      return;
+    }
+    activeDraftIdRef.current = draftSessionId;
+    syncGenerationRef.current++;
+    if (syncTimeoutRef.current != null) {
+      clearTimeout(syncTimeoutRef.current);
+      syncTimeoutRef.current = undefined;
+    }
+    pendingFlushRef.current = undefined;
+    inFlightFlushRef.current = undefined;
+    lastUpdatedAtRef.current = undefined;
+    setIsSpecSyncing(false);
+    // Dirty edits made against a previous draft must not be replayed onto
+    // the new one. Keep them only for lazy creation (undefined -> id).
+    if (previousDraftId != null) {
+      localDirtyRef.current = false;
+    }
+  }, [draftSessionId]);
+
+  useEffect(() => {
+    if (draftBridge == null) {
+      return;
+    }
+    if (draftSessionId == null) {
+      loadedDraftIdRef.current = undefined;
+      setAgentSpec(defaultAgentSpec);
+      localDirtyRef.current = false;
+      setSpecError(null);
+      setIsSpecLoading(false);
+      return;
+    }
+
+    if (loadedDraftIdRef.current === draftSessionId) {
+      return;
+    }
+
+    const abortController = new AbortController();
+    setIsSpecLoading(true);
+    void (async () => {
+      try {
+        const loaded = await draftBridge.getDraftAgentSpec(draftSessionId);
+        if (abortController.signal.aborted) {
+          return;
+        }
+        loadedDraftIdRef.current = draftSessionId;
+
+        if (localDirtyRef.current) {
+          scheduleSpecSyncRef.current(draftSessionId, agentSpecRef.current);
+          localDirtyRef.current = false;
+          setSpecError(null);
+          setIsSpecLoading(false);
+          return;
+        }
+
+        setAgentSpec(loaded);
+        setSpecError(null);
+        setIsSpecLoading(false);
+      } catch (error) {
+        if (!abortController.signal.aborted) {
+          onError?.(error);
+          setSpecError(error);
+          setIsSpecLoading(false);
+        }
+      }
+    })();
+
+    return () => {
+      abortController.abort();
+      // The cancelled load can no longer clear the flag itself. Releasing it
+      // here keeps it from sticking when the next run early-returns on an
+      // already-loaded draft; a run that starts a fresh load re-raises it in
+      // the same commit.
+      setIsSpecLoading(false);
+    };
+  }, [defaultAgentSpec, draftBridge, draftSessionId, enabled, onError]);
+
+  const flushSpecSync = useCallback(
+    async (draftId: string, spec: AgentSpec, generation: number) => {
+      if (draftBridge == null) {
+        return;
+      }
+      setIsSpecSyncing(true);
+      try {
+        const updatedAt = await draftBridge.syncAgentSpec(draftId, spec);
+        if (generation !== syncGenerationRef.current) {
+          return;
+        }
+        lastUpdatedAtRef.current = updatedAt || new Date().toISOString();
+        setSpecError(null);
+        onAgentSpecChange?.(spec);
+      } catch (error) {
+        if (generation === syncGenerationRef.current) {
+          setSpecError(error);
+          onError?.(error);
+        }
+      } finally {
+        if (generation === syncGenerationRef.current) {
+          setIsSpecSyncing(false);
+        }
+      }
+    },
+    [draftBridge, onAgentSpecChange, onError],
+  );
+
+  const scheduleSpecSync = useCallback(
+    (draftId: string, spec: AgentSpec) => {
+      if (syncTimeoutRef.current != null) {
+        clearTimeout(syncTimeoutRef.current);
+      }
+      const generation = ++syncGenerationRef.current;
+      const flush = () => {
+        pendingFlushRef.current = undefined;
+        syncTimeoutRef.current = undefined;
+        const promise = flushSpecSync(draftId, spec, generation).finally(() => {
+          if (inFlightFlushRef.current === promise) {
+            inFlightFlushRef.current = undefined;
+          }
+        });
+        inFlightFlushRef.current = promise;
+        return promise;
+      };
+      pendingFlushRef.current = flush;
+      syncTimeoutRef.current = setTimeout(() => {
+        void flush();
+      }, SPEC_SYNC_DEBOUNCE_MS);
+    },
+    [flushSpecSync],
+  );
+
+  const scheduleSpecSyncRef = useRef(scheduleSpecSync);
+  scheduleSpecSyncRef.current = scheduleSpecSync;
+
+  const flushPendingSpecSyncNow = useCallback(async () => {
+    if (syncTimeoutRef.current != null) {
+      clearTimeout(syncTimeoutRef.current);
+      syncTimeoutRef.current = undefined;
+    }
+    const pending = pendingFlushRef.current;
+    if (pending != null) {
+      pendingFlushRef.current = undefined;
+      await pending();
+      return;
+    }
+    if (inFlightFlushRef.current != null) {
+      await inFlightFlushRef.current;
+    }
+  }, []);
+
+  const adoptAgentSpec = useCallback(
+    ({ agentSpec: persistedSpec, updatedAt }: { agentSpec: AgentSpec; updatedAt?: string }) => {
+      if (syncTimeoutRef.current != null) {
+        clearTimeout(syncTimeoutRef.current);
+        syncTimeoutRef.current = undefined;
+      }
+      syncGenerationRef.current++;
+      pendingFlushRef.current = undefined;
+      inFlightFlushRef.current = undefined;
+      localDirtyRef.current = false;
+      // Keep the last successful sync timestamp when the caller omits
+      // updatedAt (e.g. a save that did not return sessionUpdatedAt).
+      if (updatedAt !== undefined) {
+        lastUpdatedAtRef.current = updatedAt;
+      }
+      agentSpecRef.current = persistedSpec;
+      setAgentSpec(persistedSpec);
+      setSpecError(null);
+      setIsSpecSyncing(false);
+    },
+    [],
+  );
+
+  const takeTurnHeaderTimestamp = useCallback(async () => {
+    await flushPendingSpecSyncNow();
+    const updatedAt = lastUpdatedAtRef.current;
+    lastUpdatedAtRef.current = undefined;
+    return updatedAt;
+  }, [flushPendingSpecSyncNow]);
+
+  useEffect(
+    () => () => {
+      if (syncTimeoutRef.current != null) {
+        clearTimeout(syncTimeoutRef.current);
+      }
+    },
+    [],
+  );
+
+  const updateAgentSpec = useCallback(
+    (update: AgentSpecUpdate) => {
+      if (draftBridge == null) {
+        return;
+      }
+      const next = mergeAgentSpec(agentSpecRef.current, update);
+      setAgentSpec(next);
+      localDirtyRef.current = true;
+      if (draftSessionId != null) {
+        scheduleSpecSync(draftSessionId, next);
+      }
+    },
+    [draftBridge, draftSessionId, enabled, scheduleSpecSync],
+  );
+
+  return useMemo(
+    () => ({
+      agentSpec: enabled ? agentSpec : null,
+      draftSessionId: enabled ? draftSessionId : undefined,
+      isSpecLoading: enabled ? isSpecLoading : false,
+      isSpecSyncing: enabled ? isSpecSyncing : false,
+      specError: enabled ? specError : null,
+      updateAgentSpec,
+      flushAgentSpec: flushPendingSpecSyncNow,
+      adoptAgentSpec,
+      takeTurnHeaderTimestamp,
+    }),
+    [
+      agentSpec,
+      draftSessionId,
+      enabled,
+      isSpecLoading,
+      isSpecSyncing,
+      specError,
+      flushPendingSpecSyncNow,
+      adoptAgentSpec,
+      takeTurnHeaderTimestamp,
+      updateAgentSpec,
+    ],
+  );
+}

package/src/extractTurnUserText.ts

@@ -0,0 +1,23 @@
+import type { Turn } from './server/index.js';
+
+export function extractTurnUserText(input: Turn['input']): string | undefined {
+  const parts: string[] = [];
+  let hasUserMessage = false;
+  for (const item of input ?? []) {
+    if (item.type !== 'user.message') {
+      continue;
+    }
+    hasUserMessage = true;
+    const { content } = item;
+    if (typeof content === 'string') {
+      parts.push(content);
+      continue;
+    }
+    for (const part of content) {
+      if (part.type === 'text') {
+        parts.push(part.text);
+      }
+    }
+  }
+  return hasUserMessage ? parts.join('\n').trim() : undefined;
+}