@convex-dev/agent 0.5.0-alpha.1 → 0.6.0-alpha.1

package/src/react/SmoothText.tsx

@@ -0,0 +1,9 @@
+import { useSmoothText, type SmoothTextOptions } from "./useSmoothText.js";
+
+export function SmoothText({
+  text,
+  ...options
+}: { text: string } & SmoothTextOptions) {
+  const [visibleText] = useSmoothText(text, options);
+  return visibleText;
+}

package/src/react/index.ts

@@ -1,237 +1,17 @@
 "use client";
-import { omit, type ErrorMessage } from "convex-helpers";
-import {
-  type PaginatedQueryArgs,
-  type UsePaginatedQueryResult,
-  useQuery,
-} from "convex/react";
-import { usePaginatedQuery } from "convex-helpers/react";
-import type { FunctionArgs } from "convex/server";
-import { useMemo, useRef, useState } from "react";
-import type { MessageDoc } from "../client/index.js";
-import type { SyncStreamsReturnValue } from "../client/types.js";
-import type { StreamArgs } from "../validators.js";
-import type { UIMessage } from "./toUIMessages.js";
-import { toUIMessages } from "./toUIMessages.js";
-import { mergeDeltas } from "./deltas.js";
-import type {
-  ThreadQuery,
-  ThreadStreamQuery,
-  ThreadMessagesArgs,
-  ThreadMessagesResult,
-} from "./types.js";
+import { useMemo, useState } from "react";
+export { toUIMessages, type UIMessage } from "../UIMessages.js";
 
 export { optimisticallySendMessage } from "./optimisticallySendMessage.js";
 export { useSmoothText } from "./useSmoothText.js";
-export { toUIMessages, type UIMessage, type ThreadQuery };
-
-/**
- * A hook that fetches messages from a thread.
- *
- * This hook is a wrapper around `usePaginatedQuery` and `useStreamingThreadMessages`.
- * It will fetch both full messages and streaming messages, and merge them together.
- *
- * The query must take as arguments `{ threadId, paginationOpts }` and return a
- * pagination result of objects that extend `MessageDoc`.
- *
- * For streaming, it should look like this:
- * ```ts
- * export const listThreadMessages = query({
- *   args: {
- *     threadId: v.string(),
- *     paginationOpts: paginationOptsValidator,
- *     streamArgs: vStreamArgs,
- *     ... other arguments you want
- *   },
- *   handler: async (ctx, { threadId, paginationOpts, streamArgs }) => {
- *     // await authorizeThreadAccess(ctx, threadId);
- *     const paginated = await agent.listMessages(ctx, { threadId, paginationOpts });
- *     const streams = await agent.syncStreams(ctx, { threadId, streamArgs });
- *     // Here you could filter out / modify the documents & stream deltas.
- *     return { ...paginated, streams };
- *   },
- * });
- * ```
- *
- * Then the hook can be used like this:
- * ```ts
- * const messages = useThreadMessages(
- *   api.myModule.listThreadMessages,
- *   { threadId },
- *   { initialNumItems: 10, stream: true }
- * );
- * ```
- *
- * @param query The query to use to fetch messages.
- * It must take as arguments `{ threadId, paginationOpts }` and return a
- * pagination result of objects that extend `MessageDoc`.
- * To support streaming, it must also take in `streamArgs: vStreamArgs` and
- * return a `streams` object returned from `agent.syncStreams`.
- * @param args The arguments to pass to the query other than `paginationOpts`
- * and `streamArgs`. So `{ threadId }` at minimum, plus any other arguments that
- * you want to pass to the query.
- * @param options The options for the query. Similar to usePaginatedQuery.
- * To enable streaming, pass `stream: true`.
- * @returns The messages. If stream is true, it will return a list of messages
- *   that includes both full messages and streaming messages.
- */
-export function useThreadMessages<
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  Query extends ThreadQuery<any, any>,
->(
-  query: Query,
-  args: ThreadMessagesArgs<Query> | "skip",
-  options: {
-    initialNumItems: number;
-    stream?: Query extends ThreadStreamQuery
-      ? boolean
-      : ErrorMessage<"To enable streaming, your query must take in streamArgs: vStreamArgs and return a streams object returned from agent.syncStreams. See docs.">;
-  },
-): UsePaginatedQueryResult<
-  ThreadMessagesResult<Query> & { streaming?: boolean }
-> {
-  // These are full messages
-  const paginated = usePaginatedQuery(
-    query,
-    args as PaginatedQueryArgs<Query> | "skip",
-    { initialNumItems: options.initialNumItems },
-  );
-
-  // These are streaming messages that will not include full messages.
-  const streamMessages = useStreamingThreadMessages(
-    query as ThreadStreamQuery<
-      ThreadMessagesArgs<Query>,
-      ThreadMessagesResult<Query>
-    >,
-    !options.stream ||
-      args === "skip" ||
-      paginated.status === "LoadingFirstPage"
-      ? "skip"
-      : { ...args, startOrder: paginated.results.at(-1)?.order },
-  );
-
-  const merged = useMemo(() => {
-    const streamListMessages =
-      streamMessages?.map((m) => ({
-        ...m,
-        streaming: !m.status || m.status === "pending",
-      })) ?? [];
-    return {
-      ...paginated,
-      results: paginated.results
-        .map((m) => ({ ...m, streaming: false }))
-        .concat(streamListMessages)
-        .sort((a, b) =>
-          a.order === b.order ? a.stepOrder - b.stepOrder : a.order - b.order,
-        )
-        // They shouldn't overlap, but check for duplicates just in case.
-        .filter(
-          (m, i, arr) =>
-            !arr[i - 1] ||
-            m.order !== arr[i - 1].order ||
-            m.stepOrder !== arr[i - 1].stepOrder,
-        ),
-    };
-  }, [paginated, streamMessages]);
-
-  return merged as ThreadMessagesResult<Query>;
-}
-
-/**
- * A hook that fetches streaming messages from a thread.
- * This ONLY returns streaming messages. To get both, use `useThreadMessages`.
- *
- * @param query The query to use to fetch messages.
- * It must take as arguments `{ threadId, paginationOpts, streamArgs }` and
- * return a `streams` object returned from `agent.syncStreams`.
- * @param args The arguments to pass to the query other than `paginationOpts`
- * and `streamArgs`. So `{ threadId }` at minimum, plus any other arguments that
- * you want to pass to the query.
- * @returns The streaming messages.
- */
-export function useStreamingThreadMessages<
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  Query extends ThreadStreamQuery<any, any>,
->(
-  query: Query,
-  args: (ThreadMessagesArgs<Query> & { startOrder?: number }) | "skip",
-): Array<ThreadMessagesResult<Query>> | undefined {
-  // Invariant: streamMessages[streamId] is comprised of all deltas up to the
-  // cursor. There can be multiple messages in the same stream, e.g. for tool
-  // calls.
-  const [streams, setStreams] = useState<
-    Array<{ streamId: string; cursor: number; messages: MessageDoc[] }>
-  >([]);
-  const startOrderRef = useRef<number>(0);
-  const queryArgs = args === "skip" ? args : omit(args, ["startOrder"]);
-  if (args !== "skip" && !startOrderRef.current && args.startOrder) {
-    startOrderRef.current = args.startOrder;
-  }
-  // Get all the active streams
-  const streamList = useQuery(
-    query,
-    queryArgs === "skip"
-      ? queryArgs
-      : ({
-          ...queryArgs,
-          paginationOpts: { cursor: null, numItems: 0 },
-          streamArgs: {
-            kind: "list",
-            startOrder: startOrderRef.current,
-          } as StreamArgs,
-        } as FunctionArgs<Query>),
-  ) as
-    | { streams: Extract<SyncStreamsReturnValue, { kind: "list" }> }
-    | undefined;
-  // Get the cursors for all the active streams
-  const cursors = useMemo(() => {
-    if (!streamList?.streams) return [];
-    if (streamList.streams.kind !== "list") {
-      throw new Error("Expected list streams");
-    }
-    return streamList.streams.messages.map(({ streamId }) => {
-      const stream = streams.find((s) => s.streamId === streamId);
-      const cursor = stream?.cursor ?? 0;
-      return { streamId, cursor };
-    });
-  }, [streamList, streams]);
-  // Get the deltas for all the active streams, if any.
-  const cursorQuery = useQuery(
-    query,
-    queryArgs === "skip" || !streamList
-      ? ("skip" as const)
-      : ({
-          ...queryArgs,
-          paginationOpts: { cursor: null, numItems: 0 },
-          streamArgs: { kind: "deltas", cursors } as StreamArgs,
-        } as FunctionArgs<Query>),
-  ) as
-    | { streams: Extract<SyncStreamsReturnValue, { kind: "deltas" }> }
-    | undefined;
-  // Merge any deltas into the streamChunks, keeping it unmodified if unchanged.
-  const threadId = args === "skip" ? undefined : args.threadId;
-  const [messages, newStreams, changed] = useMemo(() => {
-    if (!threadId) return [undefined, [], false];
-    if (!streamList) return [undefined, [], false];
-    if (cursorQuery && cursorQuery.streams?.kind !== "deltas") {
-      throw new Error("Expected deltas streams");
-    }
-    return mergeDeltas(
-      threadId,
-      streamList.streams.messages,
-      streams,
-      cursorQuery?.streams?.deltas ?? [],
-    );
-  }, [threadId, cursorQuery, streams, streamList]);
-  // Now assemble the chunks into messages
-  if (!threadId) {
-    return undefined;
-  }
-  if (changed) {
-    setStreams(newStreams);
-  }
-  return messages as ThreadMessagesResult<Query>[] | undefined;
-}
+export { SmoothText } from "./SmoothText.js";
+export {
+  type ThreadMessagesQuery,
+  useThreadMessages,
+  useStreamingThreadMessages,
+} from "./useThreadMessages.js";
+export { type UIMessagesQuery, useUIMessages } from "./useUIMessages.js";
+export { useStreamingUIMessages } from "./useStreamingUIMessages.js";
 
 /**
  * @deprecated use useThreadMessages or useStreamingThreadMessages instead

package/src/react/optimisticallySendMessage.ts

@@ -1,10 +1,52 @@
 import { insertAtTop } from "convex/react";
-import type { MessageDoc } from "../client/index.js";
+import type { MessageDoc, StreamArgs } from "../validators.js";
 import type { OptimisticLocalStore } from "convex/browser";
-import type { ThreadQuery } from "./types.js";
+import type { UIMessage } from "../UIMessages.js";
+import type {
+  FunctionReference,
+  PaginationOptions,
+  PaginationResult,
+} from "convex/server";
+import type { SyncStreamsReturnValue } from "@convex-dev/agent";
 
+/**
+ * Adds a sent message to the end of a list of messages, so it shows up until
+ * the message is saved on the server and arrives in the query.
+ * It generates a message with fields that match both MessageDoc and UIMessage,
+ * for convenience. It will not include any other fields you might have in your
+ * regular query, however.
+ *
+ * @param query The query used to fetch messages, typically with
+ * useThreadMessages or useUIMessages.
+ * @returns A function that can be used to optimistically send a message.
+ * If your mutation takes different arguments than { threadId, prompt }, you can
+ * use it as a helper function in your optimistic update:
+ * ```ts
+ * const sendMessage = useMutation(
+ *   api.chatStreaming.streamStoryAsynchronously,
+ * ).withOptimisticUpdate(
+ *   (store, args) => {
+ *     optimisticallySendMessage(api.chatStreaming.listThreadMessages)(store, {
+ *       threadId:
+ *       prompt: whatever you would have passed to the mutation,
+ *     })
+ *   }
+ * );
+ * ```
+ */
 export function optimisticallySendMessage(
-  query: ThreadQuery<unknown, MessageDoc>,
+  query: FunctionReference<
+    "query",
+    "public",
+    {
+      threadId: string;
+      paginationOpts: PaginationOptions;
+      streamArgs?: StreamArgs;
+    },
+    PaginationResult<MessageDoc | UIMessage> & {
+      streams?: SyncStreamsReturnValue;
+    }
+  >,
 ): (
   store: OptimisticLocalStore,
   args: { threadId: string; prompt: string },
@@ -12,33 +54,34 @@ export function optimisticallySendMessage(
   return (store, args) => {
     const queries = store.getAllQueries(query);
     let maxOrder = -1;
-    let maxStepOrder = 0;
     for (const q of queries) {
       if (q.args?.threadId !== args.threadId) continue;
       if (q.args.streamArgs) continue;
       for (const m of q.value?.page ?? []) {
         maxOrder = Math.max(maxOrder, m.order);
-        maxStepOrder = Math.max(maxStepOrder, m.stepOrder);
       }
     }
     const order = maxOrder + 1;
     const stepOrder = 0;
+    const id = randomUUID();
+    const { prompt, ...rest } = args;
     insertAtTop({
       paginatedQuery: query,
       argsToMatch: { threadId: args.threadId, streamArgs: undefined },
       item: {
+        ...rest,
         _creationTime: Date.now(),
-        _id: randomUUID(),
+        _id: id,
+        id,
+        key: `${args.threadId}-${order}-${stepOrder}`,
         order,
         stepOrder,
         status: "pending",
-        threadId: args.threadId,
         tool: false,
-        message: {
-          role: "user",
-          content: args.prompt,
-        },
-        text: args.prompt,
+        message: { role: "user", content: prompt },
+        parts: [{ type: "text", text: prompt }],
+        role: "user",
+        text: prompt,
       },
       localQueryStore: store,
     });

package/src/react/types.ts

@@ -1,52 +1,19 @@
 import type { BetterOmit, Expand } from "convex-helpers";
-import type {
-  FunctionArgs,
-  FunctionReference,
-  PaginationOptions,
-  PaginationResult,
-} from "convex/server";
-import type { MessageDoc } from "../client/index.js";
+import type { FunctionArgs, FunctionReference } from "convex/server";
 import type { SyncStreamsReturnValue } from "../client/types.js";
 import type { StreamArgs } from "../validators.js";
 
-export type ThreadQuery<
-  Args = unknown,
-  M extends MessageDoc = MessageDoc,
-> = FunctionReference<
+export type StreamQuery<Args = Record<string, unknown>> = FunctionReference<
   "query",
   "public",
   {
     threadId: string;
-    paginationOpts: PaginationOptions;
-    // TODO: will this allow passing a function that doesn't have this param?
-    /**
-     * If { stream: true } is passed, it will also query for stream deltas.
-     * In order for this to work, the query must take as an argument streamArgs.
-     */
-    streamArgs?: StreamArgs;
-  } & Args,
-  PaginationResult<M> & { streams?: SyncStreamsReturnValue }
->;
-
-export type ThreadStreamQuery<
-  Args = Record<string, unknown>,
-  M extends MessageDoc = MessageDoc,
-> = FunctionReference<
-  "query",
-  "public",
-  {
-    threadId: string;
-    paginationOpts: PaginationOptions;
     streamArgs?: StreamArgs; // required for stream query
   } & Args,
-  PaginationResult<M> & { streams: SyncStreamsReturnValue }
+  { streams: SyncStreamsReturnValue }
 >;
 
-export type ThreadMessagesArgs<Query extends ThreadQuery<unknown, MessageDoc>> =
-  Query extends ThreadQuery<unknown, MessageDoc>
-    ? Expand<BetterOmit<FunctionArgs<Query>, "paginationOpts" | "streamArgs">>
+export type StreamQueryArgs<Query extends StreamQuery<unknown>> =
+  Query extends StreamQuery<unknown>
+    ? Expand<BetterOmit<FunctionArgs<Query>, "streamArgs">>
     : never;
-
-export type ThreadMessagesResult<
-  Query extends ThreadQuery<unknown, MessageDoc>,
-> = Query extends ThreadQuery<unknown, infer M> ? M : never;

package/src/react/useDeltaStreams.ts

@@ -0,0 +1,160 @@
+"use client";
+
+import type { StreamQuery, StreamQueryArgs } from "./types.js";
+import type { SyncStreamsReturnValue } from "../client/types.js";
+import type { FunctionArgs } from "convex/server";
+import type { StreamArgs, StreamDelta, StreamMessage } from "../validators.js";
+import { sorted } from "../shared.js";
+import { useQuery } from "convex/react";
+import { useState } from "react";
+import { assert } from "convex-helpers";
+
+export function useDeltaStreams<
+  Query extends StreamQuery<any> = StreamQuery<object>,
+>(
+  query: Query,
+  args: StreamQueryArgs<Query> | "skip",
+  options?: {
+    startOrder?: number;
+    skipStreamIds?: string[];
+  },
+): { streamMessage: StreamMessage; deltas: StreamDelta[] }[] | undefined {
+  // We hold onto and modify state directly to avoid re-running unnecessarily.
+  const [state] = useState<{
+    startOrder: number;
+    threadId: string | undefined;
+    deltaStreams:
+      | Array<{
+          streamMessage: StreamMessage;
+          deltas: StreamDelta[];
+        }>
+      | undefined;
+  }>({
+    startOrder: options?.startOrder ?? 0,
+    deltaStreams: undefined,
+    threadId: args === "skip" ? undefined : args.threadId,
+  });
+  const [cursors, setCursors] = useState<Record<string, number>>({});
+  if (args !== "skip" && state.threadId !== args.threadId) {
+    state.threadId = args.threadId;
+    state.deltaStreams = undefined;
+    state.startOrder = options?.startOrder ?? 0;
+    setCursors({});
+  }
+  if (
+    state.deltaStreams?.length ||
+    (options?.startOrder && options.startOrder < state.startOrder)
+  ) {
+    const cacheFriendlyStartOrder = options?.startOrder
+      ? // round down to the nearest 10 for some cache benefits
+        options.startOrder - (options.startOrder % 10)
+      : 0;
+    if (cacheFriendlyStartOrder !== state.startOrder) {
+      state.startOrder = cacheFriendlyStartOrder;
+    }
+  }
+
+  // Get all the active streams
+  const streamList = useQuery(
+    query,
+    args === "skip"
+      ? args
+      : ({
+          ...args,
+          streamArgs: {
+            kind: "list",
+            startOrder: state.startOrder,
+          } as StreamArgs,
+        } as FunctionArgs<Query>),
+  ) as
+    | { streams: Extract<SyncStreamsReturnValue, { kind: "list" }> }
+    | undefined;
+
+  const streamMessages =
+    args === "skip"
+      ? undefined
+      : !streamList
+        ? state.deltaStreams?.map(({ streamMessage }) => streamMessage)
+        : sorted(
+            streamList.streams.messages.filter(
+              ({ streamId, order }) =>
+                !options?.skipStreamIds?.includes(streamId) &&
+                (!options?.startOrder || order >= options.startOrder),
+            ),
+          );
+
+  // When no active streams remain, clear the stale state so we stop
+  // returning old streaming UIMessages.
+  if (streamMessages !== undefined && streamMessages.length === 0) {
+    state.deltaStreams = undefined;
+  }
+
+  // Get the deltas for all the active streams, if any.
+  const cursorQuery = useQuery(
+    query,
+    args === "skip" || !streamMessages?.length
+      ? ("skip" as const)
+      : ({
+          ...args,
+          streamArgs: {
+            kind: "deltas",
+            cursors: streamMessages.map(({ streamId }) => ({
+              streamId,
+              cursor: cursors[streamId] ?? 0,
+            })),
+          } as StreamArgs,
+        } as FunctionArgs<Query>),
+  ) as
+    | { streams: Extract<SyncStreamsReturnValue, { kind: "deltas" }> }
+    | undefined;
+
+  const newDeltas = cursorQuery?.streams.deltas;
+  if (newDeltas?.length && streamMessages) {
+    const newDeltasByStreamId = new Map<string, StreamDelta[]>();
+    for (const delta of newDeltas) {
+      const oldCursor = cursors[delta.streamId];
+      if (oldCursor && delta.start < oldCursor) continue;
+      const existing = newDeltasByStreamId.get(delta.streamId);
+      if (existing) {
+        const previousEnd = existing.at(-1)!.end;
+        assert(
+          previousEnd === delta.start,
+          `Gap found in deltas for ${delta.streamId} jumping to ${delta.start} from ${previousEnd}`,
+        );
+        existing.push(delta);
+      } else {
+        assert(
+          !oldCursor || oldCursor === delta.start,
+          `Gap found - first delta after ${oldCursor} is ${delta.start} for stream ${delta.streamId}`,
+        );
+        newDeltasByStreamId.set(delta.streamId, [delta]);
+      }
+    }
+    const newCursors: Record<string, number> = {};
+    for (const { streamId } of streamMessages) {
+      const cursor =
+        newDeltasByStreamId.get(streamId)?.at(-1)?.end ?? cursors[streamId];
+      if (cursor !== undefined) {
+        newCursors[streamId] = cursor;
+      }
+    }
+    setCursors(newCursors);
+
+    // we defensively create a new object so object identity matches contents
+    state.deltaStreams = streamMessages.map((streamMessage) => {
+      const streamId = streamMessage.streamId;
+      const old = state.deltaStreams?.find(
+        (ds) => ds.streamMessage.streamId === streamId,
+      );
+      const newDeltas = newDeltasByStreamId.get(streamId);
+      if (!newDeltas && streamMessage === old?.streamMessage) {
+        return old;
+      }
+      return {
+        streamMessage,
+        deltas: [...(old?.deltas ?? []), ...(newDeltas ?? [])],
+      };
+    });
+  }
+  return state.deltaStreams;
+}