@assistant-ui/react 0.7.23 → 0.7.25

package/src/runtimes/remote-thread-list/RemoteThreadListThreadListRuntimeCore.tsx

@@ -0,0 +1,486 @@
+import { ThreadListRuntimeCore } from "../core/ThreadListRuntimeCore";
+import { generateId } from "../../internal";
+import { RemoteThreadListAdapter } from "./types";
+import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
+import { BaseSubscribable } from "./BaseSubscribable";
+import { EMPTY_THREAD_CORE } from "./EMPTY_THREAD_CORE";
+import { OptimisticState } from "./OptimisticState";
+import { FC, PropsWithChildren, useEffect, useId } from "react";
+import { create } from "zustand";
+import { CloudInitializeResponse } from "./cloud/CloudContext";
+
+type RemoteThreadData =
+  | {
+      readonly threadId: string;
+      readonly remoteId?: undefined;
+      readonly externalId?: undefined;
+      readonly status: "new";
+      readonly title: undefined;
+    }
+  | {
+      readonly threadId: string;
+      readonly remoteId: string;
+      readonly externalId: string | undefined;
+      readonly status: "regular" | "archived";
+      readonly title?: string | undefined;
+    };
+
+const DEFAULT_RENDER_COMPONENT: FC<PropsWithChildren> = ({ children }) => {
+  return children;
+};
+
+type THREAD_MAPPING_ID = string & { __brand: "THREAD_MAPPING_ID" };
+function createThreadMappingId(id: string): THREAD_MAPPING_ID {
+  return id as THREAD_MAPPING_ID;
+}
+
+type RemoteThreadState = {
+  readonly isLoading: boolean;
+  readonly newThreadId: string | undefined;
+  readonly threadIds: readonly string[];
+  readonly archivedThreadIds: readonly string[];
+  readonly threadIdMap: Readonly<Record<string, THREAD_MAPPING_ID>>;
+  readonly threadData: Readonly<Record<THREAD_MAPPING_ID, RemoteThreadData>>;
+};
+
+const getThreadData = (
+  state: RemoteThreadState,
+  threadIdOrRemoteId: string,
+) => {
+  const idx = state.threadIdMap[threadIdOrRemoteId];
+  if (idx === undefined) return undefined;
+  return state.threadData[idx];
+};
+
+const updateStatusReducer = (
+  state: RemoteThreadState,
+  threadIdOrRemoteId: string,
+  newStatus: "regular" | "archived" | "deleted",
+) => {
+  const data = getThreadData(state, threadIdOrRemoteId);
+  if (!data) return state;
+
+  const { threadId, remoteId, status: lastStatus } = data;
+  if (lastStatus === newStatus) return state;
+
+  const newState = { ...state };
+
+  // lastStatus
+  switch (lastStatus) {
+    case "new":
+      newState.newThreadId = undefined;
+      break;
+    case "regular":
+      newState.threadIds = newState.threadIds.filter((t) => t !== threadId);
+      break;
+    case "archived":
+      newState.archivedThreadIds = newState.archivedThreadIds.filter(
+        (t) => t !== threadId,
+      );
+      break;
+
+    default: {
+      const _exhaustiveCheck: never = lastStatus;
+      throw new Error(`Unsupported state: ${_exhaustiveCheck}`);
+    }
+  }
+
+  // newStatus
+  switch (newStatus) {
+    case "regular":
+      newState.threadIds = [...newState.threadIds, threadId];
+      break;
+
+    case "archived":
+      newState.archivedThreadIds = [...newState.archivedThreadIds, threadId];
+      break;
+
+    case "deleted":
+      newState.threadData = Object.fromEntries(
+        Object.entries(newState.threadData).filter(([key]) => key !== threadId),
+      );
+      newState.threadIdMap = Object.fromEntries(
+        Object.entries(newState.threadIdMap).filter(
+          ([key]) => key !== threadId && key !== remoteId,
+        ),
+      );
+      break;
+
+    default: {
+      const _exhaustiveCheck: never = newStatus;
+      throw new Error(`Unsupported state: ${_exhaustiveCheck}`);
+    }
+  }
+
+  if (newStatus !== "deleted") {
+    newState.threadData = {
+      ...newState.threadData,
+      [threadId]: {
+        ...data,
+        status: newStatus,
+      },
+    };
+  }
+
+  return newState;
+};
+
+export class RemoteThreadListThreadListRuntimeCore
+  extends BaseSubscribable
+  implements ThreadListRuntimeCore
+{
+  private _adapter!: RemoteThreadListAdapter;
+  private _disposeOldAdapter?: () => void;
+  private readonly _hookManager: RemoteThreadListHookInstanceManager;
+
+  private readonly _loadThreadsPromise: Promise<void>;
+
+  private _mainThreadId!: string;
+  private readonly _state = new OptimisticState<RemoteThreadState>({
+    isLoading: false,
+    newThreadId: undefined,
+    threadIds: [],
+    archivedThreadIds: [],
+    threadIdMap: {},
+    threadData: {},
+  });
+
+  public getLoadThreadsPromise() {
+    return this._loadThreadsPromise;
+  }
+
+  constructor(adapter: RemoteThreadListAdapter) {
+    super();
+
+    this._state.subscribe(() => this._notifySubscribers());
+    this._hookManager = new RemoteThreadListHookInstanceManager(
+      adapter.runtimeHook,
+    );
+    this.useRenderComponent = create(() => ({
+      RenderComponent:
+        adapter.__internal_RenderComponent ?? DEFAULT_RENDER_COMPONENT,
+    }));
+    this.__internal_setAdapter(adapter);
+
+    this._loadThreadsPromise = this._state
+      .optimisticUpdate({
+        execute: () => adapter.list(),
+        loading: (state) => {
+          return {
+            ...state,
+            isLoading: true,
+          };
+        },
+        then: (state, l) => {
+          const newThreadIds = [];
+          const newArchivedThreadIds = [];
+          const newThreadIdMap = {} as Record<string, THREAD_MAPPING_ID>;
+          const newThreadData = {} as Record<
+            THREAD_MAPPING_ID,
+            RemoteThreadData
+          >;
+
+          for (const thread of l.threads) {
+            switch (thread.status) {
+              case "regular":
+                newThreadIds.push(thread.remoteId);
+                break;
+              case "archived":
+                newArchivedThreadIds.push(thread.remoteId);
+                break;
+              default: {
+                const _exhaustiveCheck: never = thread.status;
+                throw new Error(`Unsupported state: ${_exhaustiveCheck}`);
+              }
+            }
+
+            const mappingId = createThreadMappingId(thread.remoteId);
+            newThreadIdMap[thread.remoteId] = mappingId;
+            newThreadData[mappingId] = {
+              threadId: thread.remoteId,
+              remoteId: thread.remoteId,
+              externalId: thread.externalId,
+              status: thread.status,
+              title: thread.title,
+            };
+          }
+
+          return {
+            ...state,
+            threadIds: newThreadIds,
+            archivedThreadIds: newArchivedThreadIds,
+            threadIdMap: {
+              ...state.threadIdMap,
+              ...newThreadIdMap,
+            },
+            threadData: {
+              ...state.threadData,
+              ...newThreadData,
+            },
+          };
+        },
+      })
+      .then(() => {});
+
+    this.switchToNewThread();
+  }
+
+  private useRenderComponent;
+
+  public __internal_setAdapter(adapter: RemoteThreadListAdapter) {
+    if (this._adapter === adapter) return;
+
+    this._adapter = adapter;
+    this._disposeOldAdapter?.();
+    this._disposeOldAdapter = this._adapter.onInitialize(this._onInitialize);
+
+    const RenderComponent =
+      adapter.__internal_RenderComponent ?? DEFAULT_RENDER_COMPONENT;
+    if (
+      RenderComponent !== this.useRenderComponent.getState().RenderComponent
+    ) {
+      this.useRenderComponent.setState({ RenderComponent }, true);
+    }
+
+    this._hookManager.setRuntimeHook(adapter.runtimeHook);
+  }
+
+  public get threadIds() {
+    return this._state.value.threadIds;
+  }
+
+  public get archivedThreadIds() {
+    return this._state.value.archivedThreadIds;
+  }
+
+  public get newThreadId() {
+    return this._state.value.newThreadId;
+  }
+
+  public get mainThreadId(): string {
+    return this._mainThreadId;
+  }
+
+  public getMainThreadRuntimeCore() {
+    const result = this._hookManager.getThreadRuntimeCore(this._mainThreadId);
+    if (!result) return EMPTY_THREAD_CORE;
+    return result;
+  }
+
+  public getItemById(threadIdOrRemoteId: string) {
+    return getThreadData(this._state.value, threadIdOrRemoteId);
+  }
+
+  public async switchToThread(threadIdOrRemoteId: string): Promise<void> {
+    const data = this.getItemById(threadIdOrRemoteId);
+    if (!data) throw new Error("Thread not found");
+
+    if (this._mainThreadId === data.threadId) return;
+
+    const task = this._hookManager.startThreadRuntime(data.threadId);
+    if (this.mainThreadId !== undefined) {
+      await task;
+    } else {
+      task.then(() => this._notifySubscribers());
+    }
+
+    if (data.status === "archived") await this.unarchive(data.threadId);
+    this._mainThreadId = data.threadId;
+
+    this._notifySubscribers();
+  }
+
+  public async switchToNewThread(): Promise<void> {
+    // an initialization transaction is in progress, wait for it to settle
+    while (
+      this._state.baseValue.newThreadId !== undefined &&
+      this._state.value.newThreadId === undefined
+    ) {
+      await this._state.waitForUpdate();
+    }
+
+    const state = this._state.value;
+    let threadId: string | undefined = this._state.value.newThreadId;
+    if (threadId === undefined) {
+      do {
+        threadId = `__LOCALID_${generateId()}`;
+      } while (state.threadIdMap[threadId]);
+
+      const mappingId = createThreadMappingId(threadId);
+      this._state.update({
+        ...state,
+        newThreadId: threadId,
+        threadIdMap: {
+          ...state.threadIdMap,
+          [threadId]: mappingId,
+        },
+        threadData: {
+          ...state.threadData,
+          [threadId]: {
+            status: "new",
+            threadId,
+          },
+        },
+      });
+    }
+
+    return this.switchToThread(threadId);
+  }
+
+  private _onInitialize = async (task: Promise<CloudInitializeResponse>) => {
+    const threadId = this._state.value.newThreadId;
+    if (!threadId)
+      throw new Error(
+        "ThreadListAdapter called onInitialize before switching to new thread",
+      );
+
+    await this._state.optimisticUpdate({
+      execute: () => {
+        return task;
+      },
+      optimistic: (state) => {
+        return updateStatusReducer(state, threadId, "regular");
+      },
+      then: (state, { remoteId, externalId }) => {
+        const data = getThreadData(state, threadId);
+        if (!data) return state;
+
+        const mappingId = createThreadMappingId(remoteId);
+        return {
+          ...state,
+          threadIdMap: {
+            ...state.threadIdMap,
+            [remoteId]: mappingId,
+          },
+          threadData: {
+            ...state.threadData,
+            [threadId]: {
+              ...data,
+              remoteId,
+              externalId,
+            },
+          },
+        };
+      },
+    });
+  };
+
+  public rename(threadIdOrRemoteId: string, newTitle: string): Promise<void> {
+    const data = this.getItemById(threadIdOrRemoteId);
+    if (!data) throw new Error("Thread not found");
+    if (data.status === "new") throw new Error("Thread is not yet initialized");
+
+    return this._state.optimisticUpdate({
+      execute: () => {
+        return this._adapter.rename(data.remoteId, newTitle);
+      },
+      optimistic: (state) => {
+        const data = getThreadData(state, threadIdOrRemoteId);
+        if (!data) return state;
+
+        return {
+          ...state,
+          threadData: {
+            ...state.threadData,
+            [data.threadId]: {
+              ...data,
+              title: newTitle,
+            },
+          },
+        };
+      },
+    });
+  }
+
+  private async _ensureThreadIsNotMain(threadId: string) {
+    // if thread is main thread, switch to another thread
+    if (threadId === this._mainThreadId) {
+      const lastThreadId = this._state.value.threadIds[0];
+      if (lastThreadId) {
+        await this.switchToThread(lastThreadId);
+      } else {
+        await this.switchToNewThread();
+      }
+    }
+  }
+
+  public async archive(threadIdOrRemoteId: string) {
+    const data = this.getItemById(threadIdOrRemoteId);
+    if (!data) throw new Error("Thread not found");
+    if (data.status !== "regular")
+      throw new Error("Thread is not yet initialized or already archived");
+
+    return this._state.optimisticUpdate({
+      execute: async () => {
+        await this._ensureThreadIsNotMain(data.threadId);
+        return this._adapter.archive(data.remoteId);
+      },
+      optimistic: (state) => {
+        return updateStatusReducer(state, data.threadId, "archived");
+      },
+    });
+  }
+
+  public unarchive(threadIdOrRemoteId: string): Promise<void> {
+    const data = this.getItemById(threadIdOrRemoteId);
+    if (!data) throw new Error("Thread not found");
+    if (data.status !== "archived") throw new Error("Thread is not archived");
+
+    return this._state.optimisticUpdate({
+      execute: async () => {
+        try {
+          return await this._adapter.unarchive(data.remoteId);
+        } catch (error) {
+          await this._ensureThreadIsNotMain(data.threadId);
+          throw error;
+        }
+      },
+      optimistic: (state) => {
+        return updateStatusReducer(state, data.threadId, "regular");
+      },
+    });
+  }
+
+  public async delete(threadIdOrRemoteId: string) {
+    const data = this.getItemById(threadIdOrRemoteId);
+    if (!data) throw new Error("Thread not found");
+    if (data.status !== "regular" && data.status !== "archived")
+      throw new Error("Thread is not yet initialized");
+
+    return this._state.optimisticUpdate({
+      execute: async () => {
+        await this._ensureThreadIsNotMain(data.threadId);
+        return await this._adapter.delete(data.remoteId);
+      },
+      optimistic: (state) => {
+        return updateStatusReducer(state, data.threadId, "deleted");
+      },
+    });
+  }
+
+  private useBoundIds = create<string[]>(() => []);
+
+  public __internal_RenderComponent: FC<PropsWithChildren> = ({ children }) => {
+    const id = useId();
+    useEffect(() => {
+      this.useBoundIds.setState((s) => [...s, id], true);
+      return () => {
+        this.useBoundIds.setState((s) => s.filter((i) => i !== id), true);
+      };
+    }, []);
+
+    const boundIds = this.useBoundIds();
+    const { RenderComponent } = this.useRenderComponent();
+
+    return (
+      <RenderComponent>
+        {(boundIds.length === 0 || boundIds[0] === id) && (
+          // only render if the component is the first one mounted
+          <this._hookManager.__internal_RenderThreadRuntimes />
+        )}
+
+        {children}
+      </RenderComponent>
+    );
+  };
+}

package/src/runtimes/remote-thread-list/cloud/AssistantCloud.tsx

@@ -0,0 +1,11 @@
+import { AssistantCloudAPI, AssistantCloudConfig } from "./AssistantCloudAPI";
+import { AssistantCloudThreads } from "./AssistantCloudThreads";
+
+export class AssistantCloud {
+  public readonly threads;
+
+  constructor(config: AssistantCloudConfig) {
+    const api = new AssistantCloudAPI(config);
+    this.threads = new AssistantCloudThreads(api);
+  }
+}

package/src/runtimes/remote-thread-list/cloud/AssistantCloudAPI.tsx

@@ -0,0 +1,81 @@
+import {
+  AssistantCloudAuthStrategy,
+  AssistantCloudJWTAuthStrategy,
+  AssistantCloudAPIKeyAuthStrategy,
+} from "./AssistantCloudAuthStrategy";
+
+export type AssistantCloudConfig =
+  | {
+      baseUrl: string;
+      // TODO use baseUrl to construct the projectId
+      unstable_projectId: string;
+      authToken(): Promise<string>;
+    }
+  | {
+      apiKey: string;
+      workspaceId: string;
+    };
+
+export class AssistantCloudAPI {
+  private _tokenManager: AssistantCloudAuthStrategy;
+  private _baseUrl;
+
+  constructor(config: AssistantCloudConfig) {
+    if ("authToken" in config) {
+      this._baseUrl = config.baseUrl;
+      this._tokenManager = new AssistantCloudJWTAuthStrategy(
+        config.unstable_projectId,
+        config.authToken,
+      );
+    } else {
+      this._baseUrl = "https://api.assistant-ui.com";
+      this._tokenManager = new AssistantCloudAPIKeyAuthStrategy(
+        config.apiKey,
+        config.workspaceId,
+      );
+    }
+  }
+
+  public async makeRequest(
+    endpoint: string,
+    options: {
+      method?: "POST" | "PUT" | "DELETE" | undefined;
+      query?: Record<string, string | number | boolean> | undefined;
+      body?: object | undefined;
+    } = {},
+  ) {
+    const authHeaders = await this._tokenManager.getAuthHeaders();
+    const headers = {
+      ...authHeaders,
+      "Content-Type": "application/json",
+    };
+
+    const queryParams = new URLSearchParams();
+    if (options.query) {
+      for (const [key, value] of Object.entries(options.query)) {
+        if (value === false) continue;
+        if (value === true) {
+          queryParams.set(key, "true");
+        } else {
+          queryParams.set(key, value.toString());
+        }
+      }
+    }
+
+    const url = new URL(`${this._baseUrl}${endpoint}`);
+    url.search = queryParams.toString();
+
+    const response = await fetch(url, {
+      method: options.method ?? "GET",
+      headers,
+      body: options.body ? JSON.stringify(options.body) : null,
+    });
+
+    if (!response.ok) {
+      // TODO better error handling
+      throw new Error(`Request failed with status ${response.status}`);
+    }
+
+    return response.json();
+  }
+}

package/src/runtimes/remote-thread-list/cloud/AssistantCloudAuthStrategy.tsx

@@ -0,0 +1,95 @@
+export type AssistantCloudAuthStrategy = {
+  readonly strategy: "jwt" | "api-key";
+  getAuthHeaders(): Promise<Record<string, string>>;
+};
+
+export class AssistantCloudJWTAuthStrategy
+  implements AssistantCloudAuthStrategy
+{
+  public readonly strategy = "jwt";
+
+  private cachedToken: string | null = null;
+  private tokenExpiry: number | null = null;
+
+  #projectId;
+  #authTokenCallback;
+
+  constructor(projectId: string, authTokenCallback: () => Promise<string>) {
+    this.#projectId = projectId;
+    this.#authTokenCallback = authTokenCallback;
+  }
+
+  private getJwtExpiry(jwt: string): number {
+    try {
+      const bodyPart = jwt.split(".").at(1);
+      if (!bodyPart) {
+        throw new Error("Invalid JWT format");
+      }
+
+      // Base64Url decode
+      const payload = Buffer.from(
+        bodyPart.replace(/-/g, "+").replace(/_/g, "/"),
+        "base64",
+      ).toString();
+      const payloadObj = JSON.parse(payload);
+
+      const exp = payloadObj.exp;
+      if (!exp || typeof exp !== "number") {
+        throw new Error('JWT does not contain a valid "exp" field');
+      }
+
+      // Convert exp to milliseconds
+      return exp * 1000;
+    } catch (error) {
+      throw new Error("Unable to determine the token expiry " + error);
+    }
+  }
+
+  public async getAuthHeaders(): Promise<Record<string, string>> {
+    const currentTime = Date.now();
+
+    // Check if the cached token is valid for at least 30 seconds
+    if (
+      this.cachedToken &&
+      this.tokenExpiry &&
+      this.tokenExpiry - currentTime > 30 * 1000 // 30 seconds
+    ) {
+      return {
+        Authorization: `Bearer ${this.cachedToken}`,
+        "Aui-Project-Id": this.#projectId,
+      };
+    }
+
+    // Fetch a new token
+    const newToken = await this.#authTokenCallback();
+    const expiry = this.getJwtExpiry(newToken);
+
+    this.cachedToken = newToken;
+    this.tokenExpiry = expiry;
+
+    return {
+      Authorization: `Bearer ${newToken}`,
+      "Aui-Project-Id": this.#projectId,
+    };
+  }
+}
+export class AssistantCloudAPIKeyAuthStrategy
+  implements AssistantCloudAuthStrategy
+{
+  public readonly strategy = "api-key";
+
+  #apiKey;
+  #workspaceId;
+
+  constructor(apiKey: string, workspaceId: string) {
+    this.#apiKey = apiKey;
+    this.#workspaceId = workspaceId;
+  }
+
+  public async getAuthHeaders(): Promise<Record<string, string>> {
+    return {
+      Authorization: `Bearer ${this.#apiKey}`,
+      "Aui-Workspace-Id": this.#workspaceId,
+    };
+  }
+}