@convex-dev/agent 0.2.7 → 0.2.8-alpha.0

package/src/{react/fromUIMessages.test.ts → fromUIMessages.test.ts}

@@ -1,11 +1,12 @@
-import { describe, it, expect } from "vitest";
-import { toUIMessages } from "./toUIMessages.js";
-import { fromUIMessages } from "./fromUIMessages.js";
-import type { MessageDoc } from "../client/index.js";
-import type { UIMessage } from "./toUIMessages.js";
+import { describe, expect, it } from "vitest";
+import type { MessageDoc } from "./client/index.js";
+import type { UIMessage } from "./UIMessages.js";
+import { fromUIMessages, toUIMessages } from "./UIMessages.js";
 
 // Helper to create a base message doc
-function baseMessageDoc<T = unknown>(overrides: Partial<MessageDoc & { streaming?: boolean; metadata?: T }> = {}): MessageDoc & { streaming?: boolean; metadata?: T } {
+function baseMessageDoc<T = unknown>(
+  overrides: Partial<MessageDoc & { streaming?: boolean; metadata?: T }> = {},
+): MessageDoc & { streaming?: boolean; metadata?: T } {
   return {
     _id: "msg1",
     _creationTime: Date.now(),
@@ -31,7 +32,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].role).toBe("user");
@@ -63,7 +66,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].role).toBe("assistant");
@@ -85,7 +90,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].role).toBe("system");
@@ -123,7 +130,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].text).toBe("Here's my response.");
@@ -184,35 +193,41 @@ describe("fromUIMessages round-trip tests", () => {
         tool: true,
       }),
     ];
+    const toTest = [originalMessages, [...originalMessages].reverse()];
+    for (const messages of toTest) {
+      const uiMessages = toUIMessages(messages);
+      const backToMessageDocs = fromUIMessages(uiMessages, {
+        threadId: "thread1",
+      });
 
-    const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
-
-    // Should be grouped into single UI message
-    expect(uiMessages).toHaveLength(1);
-    expect(uiMessages[0].role).toBe("assistant");
-
-    // Check tool parts exist
-    const toolParts = uiMessages[0].parts.filter(
-      (part) => part.type === "tool-calculator",
-    );
-    expect(toolParts).toHaveLength(1);
-    expect(toolParts[0]).toMatchObject({
-      type: "tool-calculator",
-      toolCallId: "call1",
-      state: "output-available",
-      input: { operation: "add", a: 2, b: 3 },
-      output: { result: 5 },
-    });
+      // Should be grouped into single UI message
+      expect(uiMessages).toHaveLength(1);
+      const uiMessage = uiMessages[0];
+      expect(uiMessage.role).toBe("assistant");
+      expect(uiMessage.id).toBe("msg1");
+
+      // Check tool parts exist
+      const toolParts = uiMessage.parts.filter(
+        (part) => part.type === "tool-calculator",
+      );
+      expect(toolParts).toHaveLength(1);
+      expect(toolParts[0]).toMatchObject({
+        type: "tool-calculator",
+        toolCallId: "call1",
+        state: "output-available",
+        input: { operation: "add", a: 2, b: 3 },
+        output: { result: 5 },
+      });
 
-    // Should expand back to multiple message docs
-    expect(backToMessageDocs.length).toBeGreaterThanOrEqual(1);
+      // Should expand back to multiple message docs
+      expect(backToMessageDocs.length).toBeGreaterThanOrEqual(1);
 
-    // Check that tool information is preserved
-    const toolMessages = backToMessageDocs.filter((msg) => msg.tool);
-    expect(toolMessages.length).toBeGreaterThan(0);
-    expect(toolMessages[0].stepOrder).toBe(1);
-    expect(toolMessages[1].stepOrder).toBe(2);
+      // Check that tool information is preserved
+      const toolMessages = backToMessageDocs.filter((msg) => msg.tool);
+      expect(toolMessages.length).toBeGreaterThan(0);
+      expect(toolMessages[0].stepOrder).toBe(1);
+      expect(toolMessages[1].stepOrder).toBe(2);
+    }
   });
 
   it("preserves file attachments in user messages", () => {
@@ -237,7 +252,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].role).toBe("user");
@@ -298,7 +315,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
 
@@ -337,7 +356,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].metadata).toEqual(testMetadata);
@@ -360,7 +381,9 @@ describe("fromUIMessages round-trip tests", () => {
     ];
 
     const uiMessages = toUIMessages(originalMessages);
-    const backToMessageDocs = fromUIMessages("thread1", uiMessages);
+    const backToMessageDocs = fromUIMessages(uiMessages, {
+      threadId: "thread1",
+    });
 
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].status).toBe("streaming");
@@ -374,7 +397,7 @@ describe("fromUIMessages round-trip tests", () => {
 describe("fromUIMessages functionality tests", () => {
   it("handles empty messages array", () => {
     const uiMessages: UIMessage[] = [];
-    const result = fromUIMessages("thread1", uiMessages);
+    const result = fromUIMessages(uiMessages, { threadId: "thread1" });
     expect(result).toHaveLength(0);
   });
 
@@ -391,7 +414,9 @@ describe("fromUIMessages functionality tests", () => {
       parts: [{ type: "text", text: "Hello" }],
     };
 
-    const result = fromUIMessages("custom-thread-id", [uiMessage]);
+    const result = fromUIMessages([uiMessage], {
+      threadId: "custom-thread-id",
+    });
     expect(result).toHaveLength(1);
     expect(result[0].threadId).toBe("custom-thread-id");
   });
@@ -417,7 +442,7 @@ describe("fromUIMessages functionality tests", () => {
       ],
     };
 
-    const result = fromUIMessages("thread1", [toolUIMessage]);
+    const result = fromUIMessages([toolUIMessage], { threadId: "thread1" });
     expect(result.length).toBeGreaterThan(0);
 
     // Should have tool messages
@@ -446,7 +471,7 @@ describe("fromUIMessages functionality tests", () => {
       ],
     };
 
-    const result = fromUIMessages("thread1", [toolUIMessage]);
+    const result = fromUIMessages([toolUIMessage], { threadId: "thread1" });
     expect(result.length).toBeGreaterThan(0);
 
     // Should have tool messages

package/src/react/index.ts

@@ -1,259 +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, 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";
+export { toUIMessages, type UIMessage } from "../UIMessages.js";
 
 export { optimisticallySendMessage } from "./optimisticallySendMessage.js";
 export { useSmoothText } from "./useSmoothText.js";
 export { SmoothText } from "./SmoothText.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 },
-  );
-
-  let startOrder = paginated.results.at(-1)?.order ?? 0;
-  for (let i = paginated.results.length - 1; i >= 0; i--) {
-    const m = paginated.results[i];
-    if (!m.streaming && m.status === "pending") {
-      // round down to the nearest 10 for some cache benefits
-      startOrder = m.order - (m.order % 10);
-      break;
-    }
-  }
-  // 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 },
-  );
-
-  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 }))
-        // Note: this is intentionally after paginated results.
-        .concat(streamListMessages)
-        .sort((a, b) =>
-          a.order === b.order ? a.stepOrder - b.stepOrder : a.order - b.order,
-        )
-        .reduce(
-          (msgs, msg) => {
-            const last = msgs.at(-1);
-            if (!last) {
-              return [msg];
-            }
-            if (last.order !== msg.order || last.stepOrder !== msg.stepOrder) {
-              return [...msgs, msg];
-            }
-            if (
-              last.status === "pending" &&
-              (msg.streaming || msg.status !== "pending")
-            ) {
-              // Let's prefer a streaming or finalized message over a pending
-              // one.
-              return [...msgs.slice(0, -1), msg];
-            }
-            // skip the new one if the previous one (listed) was finalized
-            return msgs;
-          },
-          [] as (ThreadMessagesResult<Query> & { streaming: boolean })[],
-        ),
-    };
-  }, [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 queryArgs = args === "skip" ? args : omit(args, ["startOrder"]);
-  // Get all the active streams
-  const streamList = useQuery(
-    query,
-    queryArgs === "skip"
-      ? queryArgs
-      : ({
-          ...queryArgs,
-          paginationOpts: { cursor: null, numItems: 0 },
-          streamArgs: {
-            kind: "list",
-            startOrder: queryArgs.startOrder ?? 0,
-          } 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 {
+  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 },
@@ -20,19 +62,26 @@ export function optimisticallySendMessage(
       }
     }
     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: 0,
+        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 StreamMessagesArgs<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;