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

package/src/react/useUIMessages.test.ts

@@ -0,0 +1,255 @@
+import { describe, it, expect } from "vitest";
+import { dedupeMessages } from "./useUIMessages.js";
+
+type TestMessage = {
+  order: number;
+  stepOrder: number;
+  status: "pending" | "success" | "failed" | "streaming";
+  id: string;
+};
+
+describe("dedupeMessages", () => {
+  it("should prefer messages from messages list when streaming messages are absent", () => {
+    const messages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "success", id: "msg1" },
+      { order: 2, stepOrder: 0, status: "success", id: "msg2" },
+    ];
+    const streamMessages: TestMessage[] = [];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(2);
+    expect(result).toEqual(messages);
+  });
+
+  it("should prefer streaming messages when messages list is empty", () => {
+    const messages: TestMessage[] = [];
+    const streamMessages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "streaming", id: "stream1" },
+      { order: 2, stepOrder: 0, status: "streaming", id: "stream2" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(2);
+    expect(result).toEqual(streamMessages);
+  });
+
+  it("should prefer non-pending messages over pending messages", () => {
+    const messages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "pending", id: "pending1" },
+      { order: 2, stepOrder: 0, status: "success", id: "success1" },
+    ];
+    const streamMessages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "success", id: "finalized1" },
+      { order: 2, stepOrder: 0, status: "streaming", id: "streaming1" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(2);
+    expect(result[0]).toEqual({
+      order: 1,
+      stepOrder: 0,
+      status: "success",
+      id: "finalized1",
+    });
+    expect(result[1]).toEqual({
+      order: 2,
+      stepOrder: 0,
+      status: "success",
+      id: "success1",
+    });
+  });
+
+  it("should prefer streaming over pending messages", () => {
+    const messages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "pending", id: "pending1" },
+    ];
+    const streamMessages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "streaming", id: "streaming1" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(1);
+    expect(result[0]).toEqual({
+      order: 1,
+      stepOrder: 0,
+      status: "streaming",
+      id: "streaming1",
+    });
+  });
+
+  it("should prefer non-pending messages from messages list over streaming messages", () => {
+    const messages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "success", id: "success1" },
+    ];
+    const streamMessages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "streaming", id: "streaming1" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(1);
+    expect(result[0]).toEqual({
+      order: 1,
+      stepOrder: 0,
+      status: "success",
+      id: "success1",
+    });
+  });
+
+  it("should handle complex scenarios with multiple duplicates", () => {
+    const messages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "pending", id: "pending1" },
+      { order: 2, stepOrder: 0, status: "success", id: "success2" },
+      { order: 3, stepOrder: 0, status: "failed", id: "failed3" },
+      { order: 4, stepOrder: 0, status: "success", id: "success4" },
+    ];
+    const streamMessages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "streaming", id: "streaming1" },
+      { order: 2, stepOrder: 0, status: "streaming", id: "streaming2" },
+      { order: 5, stepOrder: 0, status: "streaming", id: "streaming5" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(5);
+    // Order 1: streaming should replace pending
+    expect(result[0]).toEqual({
+      order: 1,
+      stepOrder: 0,
+      status: "streaming",
+      id: "streaming1",
+    });
+    // Order 2: success should be kept over streaming
+    expect(result[1]).toEqual({
+      order: 2,
+      stepOrder: 0,
+      status: "success",
+      id: "success2",
+    });
+    // Order 3: failed should be kept (no streaming equivalent)
+    expect(result[2]).toEqual({
+      order: 3,
+      stepOrder: 0,
+      status: "failed",
+      id: "failed3",
+    });
+    // Order 4: success should be kept (no streaming equivalent)
+    expect(result[3]).toEqual({
+      order: 4,
+      stepOrder: 0,
+      status: "success",
+      id: "success4",
+    });
+    // Order 5: streaming should be added (no messages equivalent)
+    expect(result[4]).toEqual({
+      order: 5,
+      stepOrder: 0,
+      status: "streaming",
+      id: "streaming5",
+    });
+  });
+
+  it("should handle different stepOrder values correctly", () => {
+    const messages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "pending", id: "pending1-0" },
+      { order: 1, stepOrder: 1, status: "success", id: "success1-1" },
+    ];
+    const streamMessages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "streaming", id: "streaming1-0" },
+      { order: 1, stepOrder: 2, status: "streaming", id: "streaming1-2" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(3);
+    // Order 1, stepOrder 0: streaming should replace pending
+    expect(result[0]).toEqual({
+      order: 1,
+      stepOrder: 0,
+      status: "streaming",
+      id: "streaming1-0",
+    });
+    // Order 1, stepOrder 1: success should be kept
+    expect(result[1]).toEqual({
+      order: 1,
+      stepOrder: 1,
+      status: "success",
+      id: "success1-1",
+    });
+    // Order 1, stepOrder 2: streaming should be added
+    expect(result[2]).toEqual({
+      order: 1,
+      stepOrder: 2,
+      status: "streaming",
+      id: "streaming1-2",
+    });
+  });
+
+  it("should maintain proper sorting by order and stepOrder", () => {
+    const messages: TestMessage[] = [
+      { order: 3, stepOrder: 0, status: "success", id: "msg3" },
+      { order: 1, stepOrder: 1, status: "success", id: "msg1-1" },
+      { order: 1, stepOrder: 0, status: "pending", id: "msg1-0" },
+    ];
+    const streamMessages: TestMessage[] = [
+      { order: 2, stepOrder: 0, status: "streaming", id: "stream2" },
+      { order: 1, stepOrder: 0, status: "streaming", id: "stream1-0" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(4);
+    // Should be sorted by order, then stepOrder
+    expect(result[0]).toEqual({
+      order: 1,
+      stepOrder: 0,
+      status: "streaming",
+      id: "stream1-0",
+    });
+    expect(result[1]).toEqual({
+      order: 1,
+      stepOrder: 1,
+      status: "success",
+      id: "msg1-1",
+    });
+    expect(result[2]).toEqual({
+      order: 2,
+      stepOrder: 0,
+      status: "streaming",
+      id: "stream2",
+    });
+    expect(result[3]).toEqual({
+      order: 3,
+      stepOrder: 0,
+      status: "success",
+      id: "msg3",
+    });
+  });
+
+  it("should handle empty arrays", () => {
+    const result = dedupeMessages([], []);
+    expect(result).toEqual([]);
+  });
+
+  it("should demonstrate order dependency when messages and streamMessages have same order/stepOrder", () => {
+    // This test shows that the current implementation doesn't guarantee
+    // which message wins when they have the same order/stepOrder and status
+    const messages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "success", id: "messages-success" },
+    ];
+    const streamMessages: TestMessage[] = [
+      { order: 1, stepOrder: 0, status: "success", id: "stream-success" },
+    ];
+
+    const result = dedupeMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(1);
+    // The result depends on which array comes first in messages.concat(streamMessages)
+    // Since messages comes first, it should win when both have same status
+    expect(result[0].id).toBe("messages-success");
+  });
+});

package/src/react/useUIMessages.ts

@@ -0,0 +1,195 @@
+"use client";
+import {
+  type BetterOmit,
+  type ErrorMessage,
+  type Expand,
+} from "convex-helpers";
+import { usePaginatedQuery } from "convex-helpers/react";
+import {
+  type PaginatedQueryArgs,
+  type UsePaginatedQueryResult,
+} from "convex/react";
+import type {
+  FunctionArgs,
+  FunctionReference,
+  PaginationOptions,
+  PaginationResult,
+} from "convex/server";
+import { useMemo } from "react";
+import type { SyncStreamsReturnValue } from "../client/types.js";
+import type { StreamArgs } from "../validators.js";
+import type { StreamQuery } from "./types.js";
+import {
+  type UIMessage,
+  type UIStatus,
+  combineUIMessages,
+} from "../UIMessages.js";
+import { sorted } from "../shared.js";
+import { useStreamingUIMessages } from "./useStreamingUIMessages.js";
+
+export type UIMessageLike = {
+  order: number;
+  stepOrder: number;
+  status: UIStatus;
+  parts: UIMessage["parts"];
+  role: UIMessage["role"];
+};
+
+export type UIMessagesQuery<
+  Args = unknown,
+  M extends UIMessageLike = UIMessageLike,
+> = FunctionReference<
+  "query",
+  "public",
+  {
+    threadId: string;
+    paginationOpts: PaginationOptions;
+    /**
+     * 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 UIMessagesQueryArgs<
+  Query extends UIMessagesQuery<unknown, UIMessageLike>,
+> =
+  Query extends UIMessagesQuery<unknown, UIMessageLike>
+    ? Expand<BetterOmit<FunctionArgs<Query>, "paginationOpts" | "streamArgs">>
+    : never;
+
+export type UIMessagesQueryResult<
+  Query extends UIMessagesQuery<unknown, UIMessageLike>,
+> = Query extends UIMessagesQuery<unknown, infer M> ? M : never;
+
+/**
+ * A hook that fetches UIMessages from a thread.
+ *
+ * It's similar to useThreadMessages, for endpoints that return UIMessages.
+ * The streaming messages are materialized as UIMessages. The rest are passed
+ * through from the query.
+ *
+ * This hook is a wrapper around `usePaginatedQuery` and `useStreamingUIMessages`.
+ * 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 similar to UIMessage:
+ *
+ * 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, args) => {
+ *     // await authorizeThreadAccess(ctx, threadId);
+ *     // NOTE: listUIMessages returns UIMessages, not MessageDocs.
+ *     const paginated = await listUIMessages(ctx, components.agent, args);
+ *     const streams = await syncStreams(ctx, components.agent, args);
+ *     // Here you could filter out / modify the documents & stream deltas.
+ *     return { ...paginated, streams };
+ *   },
+ * });
+ * ```
+ *
+ * Then the hook can be used like this:
+ * ```ts
+ * const { results, status, loadMore } = useUIMessages(
+ *   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 similar to UIMessage:
+ * Required fields: (role, parts, status, order, stepOrder).
+ * To support streaming, it must also take in `streamArgs: vStreamArgs` and
+ * return a `streams` object returned from `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.
+ *   The streaming messages are materialized as UIMessages. The rest are passed
+ *   through from the query.
+ */
+export function useUIMessages<Query extends UIMessagesQuery<any, any>>(
+  query: Query,
+  args: UIMessagesQueryArgs<Query> | "skip",
+  options: {
+    initialNumItems: number;
+    stream?: Query extends StreamQuery
+      ? boolean
+      : ErrorMessage<"To enable streaming, your query must take in streamArgs: vStreamArgs and return a streams object returned from syncStreams. See docs.">;
+    skipStreamIds?: string[];
+  },
+): UsePaginatedQueryResult<UIMessagesQueryResult<Query>> {
+  // These are full messages
+  const paginated = usePaginatedQuery(
+    query,
+    args as PaginatedQueryArgs<Query> | "skip",
+    { initialNumItems: options.initialNumItems },
+  );
+
+  const startOrder = paginated.results.length
+    ? Math.min(...paginated.results.map((m) => m.order))
+    : 0;
+  // These are streaming messages that will not include full messages.
+  const streamMessages = useStreamingUIMessages(
+    query as StreamQuery<UIMessagesQueryArgs<Query>>,
+    !options.stream ||
+      args === "skip" ||
+      paginated.status === "LoadingFirstPage"
+      ? "skip"
+      : ({ ...args, paginationOpts: { cursor: null, numItems: 0 } } as any),
+    { startOrder, skipStreamIds: options.skipStreamIds },
+  );
+
+  const merged = useMemo(() => {
+    // Messages may have been split by pagination. Re-combine them here.
+    const combined = combineUIMessages(sorted(paginated.results));
+    return {
+      ...paginated,
+      results: dedupeMessages(combined, streamMessages ?? []),
+    };
+  }, [paginated, streamMessages]);
+
+  return merged as UIMessagesQueryResult<Query>;
+}
+
+export function dedupeMessages<
+  M extends {
+    order: number;
+    stepOrder: number;
+    status: UIStatus;
+  },
+>(messages: M[], streamMessages: M[]): M[] {
+  return sorted(messages.concat(streamMessages)).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" || last.status === "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 M[]);
+}

package/src/shared.ts

@@ -1,9 +1,23 @@
-import type { CoreMessage } from "ai";
-import type { Message } from "./validators.js";
+import type {
+  FilePart,
+  ImagePart,
+  ReasoningPart,
+  ToolCallPart,
+  ToolResultPart,
+  ToolApprovalRequest,
+} from "@ai-sdk/provider-utils";
+import type {
+  ModelMessage,
+  TextPart,
+  UIDataTypes,
+  UIMessagePart,
+  UITools,
+} from "ai";
+import type { Message, MessageContentParts } from "./validators.js";
 
 export const DEFAULT_RECENT_MESSAGES = 100;
 
-export function isTool(message: Message | CoreMessage) {
+export function isTool(message: Message | ModelMessage) {
   return (
     message.role === "tool" ||
     (message.role === "assistant" &&
@@ -12,26 +26,88 @@ export function isTool(message: Message | CoreMessage) {
   );
 }
 
-export function extractText(message: Message | CoreMessage) {
+export function extractText(message: Message | ModelMessage) {
   switch (message.role) {
     case "user":
       if (typeof message.content === "string") {
         return message.content;
       }
-      return message.content
-        .filter((c) => c.type === "text")
-        .map((c) => c.text)
-        .join("");
+      return joinText(message.content);
     case "assistant":
       if (typeof message.content === "string") {
         return message.content;
+      } else {
+        return joinText(message.content) || undefined;
       }
-      return message.content
-        .filter((c) => c.type === "text")
-        .map((c) => c.text)
-        .join("");
+    case "system":
+      return message.content;
+    // we don't extract text from tool messages
   }
   return undefined;
 }
 
+export function joinText(
+  parts: (
+    | UIMessagePart<UIDataTypes, UITools>
+    | TextPart
+    | ImagePart
+    | FilePart
+    | ReasoningPart
+    | ToolCallPart
+    | ToolResultPart
+    | MessageContentParts
+    | ToolApprovalRequest
+  )[],
+) {
+  return parts
+    .filter((p) => p.type === "text")
+    .map((p) => p.text)
+    .filter(Boolean)
+    .join(" ");
+}
+
+export function extractReasoning(message: Message | ModelMessage) {
+  if (typeof message.content === "string") {
+    return undefined;
+  }
+  return message.content
+    .filter((c) => c.type === "reasoning")
+    .map((c) => c.text)
+    .join(" ");
+}
+
 export const DEFAULT_MESSAGE_RANGE = { before: 2, after: 1 };
+
+export function sorted<T extends { order: number; stepOrder: number }>(
+  messages: T[],
+  order: "asc" | "desc" = "asc",
+): T[] {
+  return [...messages].sort(
+    order === "asc"
+      ? (a, b) => a.order - b.order || a.stepOrder - b.stepOrder
+      : (a, b) => b.order - a.order || b.stepOrder - a.stepOrder,
+  );
+}
+
+export type ModelOrMetadata =
+  | string
+  | ({ provider: string } & ({ modelId: string } | { model: string }));
+
+export function getModelName(embeddingModel: ModelOrMetadata): string {
+  if (typeof embeddingModel === "string") {
+    if (embeddingModel.includes("/")) {
+      return embeddingModel.split("/").slice(1).join("/");
+    }
+    return embeddingModel;
+  }
+  return "modelId" in embeddingModel
+    ? embeddingModel.modelId
+    : embeddingModel.model;
+}
+
+export function getProviderName(embeddingModel: ModelOrMetadata): string {
+  if (typeof embeddingModel === "string") {
+    return embeddingModel.split("/").at(0)!;
+  }
+  return embeddingModel.provider;
+}

package/src/test.ts

@@ -0,0 +1,18 @@
+/// <reference types="vite/client" />
+import type { TestConvex } from "convex-test";
+import type { GenericSchema, SchemaDefinition } from "convex/server";
+import schema from "./component/schema.js";
+const modules = import.meta.glob("./component/**/*.ts");
+
+/**
+ * Register the component with the test convex instance.
+ * @param t - The test convex instance, e.g. from calling `convexTest`.
+ * @param name - The name of the component, as registered in convex.config.ts.
+ */
+export function register(
+  t: TestConvex<SchemaDefinition<GenericSchema, boolean>>,
+  name: string = "agent",
+) {
+  t.registerComponent(name, schema, modules);
+}
+export default { register, schema, modules };