@convex-dev/agent 0.1.14 → 0.1.15-alpha.1

package/src/client/listMessages.ts

@@ -0,0 +1,38 @@
+import type { PaginationOptions, PaginationResult } from "convex/server";
+import type { AgentComponent, RunQueryCtx } from "./types.js";
+import type { MessageStatus } from "../validators.js";
+import type { MessageDoc } from "../component/schema.js";
+
+/**
+ * List messages from a thread.
+ * @param ctx A ctx object from a query, mutation, or action.
+ * @param component The agent component, usually `components.agent`.
+ * @param args.threadId The thread to list messages from.
+ * @param args.paginationOpts Pagination options (e.g. via usePaginatedQuery).
+ * @param args.excludeToolMessages Whether to exclude tool messages.
+ *   False by default.
+ * @param args.statuses What statuses to include. All by default.
+ * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
+ */
+export async function listMessages(
+  ctx: RunQueryCtx,
+  component: AgentComponent,
+  args: {
+    threadId: string;
+    paginationOpts: PaginationOptions;
+    excludeToolMessages?: boolean;
+    statuses?: MessageStatus[];
+  }
+): Promise<PaginationResult<MessageDoc>> {
+  if (args.paginationOpts.numItems === 0) {
+    return {
+      page: [],
+      isDone: true,
+      continueCursor: args.paginationOpts.cursor ?? "",
+    };
+  }
+  return ctx.runQuery(component.messages.listMessagesByThreadId, {
+    order: "desc",
+    ...args,
+  });
+}

package/src/client/search.ts

@@ -0,0 +1,172 @@
+import type {
+  AgentComponent,
+  ContextOptions,
+  RunActionCtx,
+  RunQueryCtx,
+} from "./types.js";
+import type { MessageDoc } from "../component/schema.js";
+import type { CoreMessage } from "ai";
+import { assert } from "convex-helpers";
+import {
+  DEFAULT_MESSAGE_RANGE,
+  DEFAULT_RECENT_MESSAGES,
+  extractText,
+} from "../shared.js";
+
+const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0;
+
+export type GetEmbedding = (text: string) => Promise<{
+  vector: number[];
+  vectorModel: string;
+  vectorScoreThreshold?: number;
+}>;
+
+/**
+ * Fetch the context messages for a thread.
+ * @param ctx Either a query, mutation, or action ctx.
+ *   If it is not an action context, you can't do text or
+ *   vector search.
+ * @param args The associated thread, user, message
+ * @returns
+ */
+export async function fetchContextMessages(
+  ctx: RunQueryCtx | RunActionCtx,
+  component: AgentComponent,
+  args: {
+    userId: string | undefined;
+    threadId: string | undefined;
+    messages: CoreMessage[];
+    /**
+     * If provided, it will search for messages up to and including this message.
+     * Note: if this is far in the past, text and vector search results may be more
+     * limited, as it's post-filtering the results.
+     */
+    upToAndIncludingMessageId?: string;
+    contextOptions: ContextOptions;
+    getEmbedding?: GetEmbedding;
+  }
+): Promise<MessageDoc[]> {
+  assert(args.userId || args.threadId, "Specify userId or threadId");
+  const opts = args.contextOptions;
+  // Fetch the latest messages from the thread
+  let included: Set<string> | undefined;
+  const contextMessages: MessageDoc[] = [];
+  if (
+    args.threadId &&
+    (opts.recentMessages !== 0 || args.upToAndIncludingMessageId)
+  ) {
+    const { page } = await ctx.runQuery(
+      component.messages.listMessagesByThreadId,
+      {
+        threadId: args.threadId,
+        excludeToolMessages: opts.excludeToolMessages,
+        paginationOpts: {
+          numItems: opts.recentMessages ?? DEFAULT_RECENT_MESSAGES,
+          cursor: null,
+        },
+        upToAndIncludingMessageId: args.upToAndIncludingMessageId,
+        order: "desc",
+        statuses: ["success"],
+      }
+    );
+    included = new Set(page.map((m) => m._id));
+    contextMessages.push(
+      // Reverse since we fetched in descending order
+      ...page.reverse()
+    );
+  }
+  if (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) {
+    const targetMessage = contextMessages.find(
+      (m) => m._id === args.upToAndIncludingMessageId
+    )?.message;
+    const messagesToSearch = targetMessage ? [targetMessage] : args.messages;
+    if (!("runAction" in ctx)) {
+      throw new Error("searchUserMessages only works in an action");
+    }
+    const lastMessage = messagesToSearch.at(-1)!;
+    assert(lastMessage, "No messages to search");
+    const text = extractText(lastMessage);
+    assert(text, `No text to search in message ${JSON.stringify(lastMessage)}`);
+    assert(
+      !args.contextOptions?.searchOptions?.vectorSearch || "runAction" in ctx,
+      "You must do vector search from an action"
+    );
+    if (opts.searchOptions?.vectorSearch && !args.getEmbedding) {
+      throw new Error(
+        "You must provide an embedding and embeddingModel to use vector search"
+      );
+    }
+    const embeddingFields = opts.searchOptions?.vectorSearch
+      ? await args.getEmbedding?.(text)
+      : undefined;
+    const searchMessages = await ctx.runAction(
+      component.messages.searchMessages,
+      {
+        searchAllMessagesForUserId: opts?.searchOtherThreads
+          ? args.userId ??
+            (args.threadId &&
+              (
+                await ctx.runQuery(component.threads.getThread, {
+                  threadId: args.threadId,
+                })
+              )?.userId)
+          : undefined,
+        threadId: args.threadId,
+        beforeMessageId: args.upToAndIncludingMessageId,
+        limit: opts.searchOptions?.limit ?? 10,
+        messageRange: {
+          ...DEFAULT_MESSAGE_RANGE,
+          ...opts.searchOptions?.messageRange,
+        },
+        text,
+        vectorScoreThreshold:
+          opts.searchOptions?.vectorScoreThreshold ??
+          DEFAULT_VECTOR_SCORE_THRESHOLD,
+        ...embeddingFields,
+      }
+    );
+    // TODO: track what messages we used for context
+    contextMessages.unshift(
+      ...searchMessages.filter((m) => !included?.has(m._id))
+    );
+  }
+  // Ensure we don't include tool messages without a corresponding tool call
+  return filterOutOrphanedToolMessages(
+    contextMessages.sort((a, b) =>
+      // Sort the raw MessageDocs by order and stepOrder
+      a.order === b.order ? a.stepOrder - b.stepOrder : a.order - b.order
+    )
+  );
+}
+
+/**
+ * Filter out tool messages that don't have both a tool call and response.
+ * @param docs The messages to filter.
+ * @returns The filtered messages.
+ */
+export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
+  const toolCallIds = new Set<string>();
+  const result: MessageDoc[] = [];
+  for (const doc of docs) {
+    if (
+      doc.message?.role === "assistant" &&
+      Array.isArray(doc.message.content)
+    ) {
+      for (const content of doc.message.content) {
+        if (content.type === "tool-call") {
+          toolCallIds.add(content.toolCallId);
+        }
+      }
+      result.push(doc);
+    } else if (doc.message?.role === "tool") {
+      if (doc.message.content.every((c) => toolCallIds.has(c.toolCallId))) {
+        result.push(doc);
+      } else {
+        console.debug("Filtering out orphaned tool message", doc);
+      }
+    } else {
+      result.push(doc);
+    }
+  }
+  return result;
+}

package/src/client/streaming.ts

@@ -6,11 +6,58 @@ import {
 } from "ai";
 import type {
   ProviderOptions,
+  StreamArgs,
   StreamDelta,
   TextStreamPart,
 } from "../validators.js";
-import type { AgentComponent, MessageDoc } from "./index.js";
-import type { RunActionCtx } from "./types.js";
+import type { MessageDoc } from "../component/schema.js";
+import type {
+  AgentComponent,
+  RunActionCtx,
+  RunQueryCtx,
+  SyncStreamsReturnValue,
+} from "./types.js";
+
+/**
+ * A function that handles fetching stream deltas, used with the React hooks
+ * `useThreadMessages` or `useStreamingThreadMessages`.
+ * @param ctx A ctx object from a query, mutation, or action.
+ * @param component The agent component, usually `components.agent`.
+ * @param args.threadId The thread to sync streams for.
+ * @param args.streamArgs The stream arguments with per-stream cursors.
+ * @returns The deltas for each stream from their existing cursor.
+ */
+export async function syncStreams(
+  ctx: RunQueryCtx,
+  component: AgentComponent,
+  args: {
+    threadId: string;
+    streamArgs: StreamArgs | undefined;
+    // By default, only streaming messages are included.
+    includeStatuses?: ("streaming" | "finished" | "aborted")[];
+  }
+): Promise<SyncStreamsReturnValue | undefined> {
+  if (!args.streamArgs) return undefined;
+  if (args.streamArgs.kind === "list") {
+    return {
+      kind: "list",
+      messages: await ctx.runQuery(component.streams.list, {
+        threadId: args.threadId,
+        startOrder: args.streamArgs.startOrder,
+        statuses: args.includeStatuses,
+      }),
+    };
+  } else {
+    return {
+      kind: "deltas",
+      deltas: await ctx.runQuery(component.streams.listDeltas, {
+        threadId: args.threadId,
+        cursors: args.streamArgs.cursors,
+      }),
+    };
+  }
+}
+
 
 export type StreamingOptions = {
   /**
@@ -96,7 +143,13 @@ export class DeltaStreamer {
     this.#nextStepOrder = (metadata.stepOrder ?? 0) + 1;
     this.abortController = new AbortController();
     if (metadata.abortSignal) {
-      metadata.abortSignal.addEventListener("abort", () => {
+      metadata.abortSignal.addEventListener("abort", async () => {
+        if (this.streamId) {
+          await this.ctx.runMutation(this.component.streams.abort, {
+            streamId: this.streamId,
+            reason: "abortSignal",
+          });
+        }
         this.abortController.abort();
       });
     }

package/src/client/types.ts

@@ -77,6 +77,10 @@ export type ContextOptions = {
      * At least one of textSearch or vectorSearch must be true.
      */
     vectorSearch?: boolean;
+    /**
+     * The score threshold for vector search. Default is 0.0.
+     */
+    vectorScoreThreshold?: number;
     /**
      * What messages around the search results to include.
      * Default: { before: 2, after: 1 }
@@ -144,7 +148,7 @@ export type UsageHandler = (
 ) => void | Promise<void>;
 
 export type RawRequestResponseHandler = (
-  ctx: RunActionCtx,
+  ctx: ActionCtx,
   args: {
     userId: string | undefined;
     threadId: string | undefined;

package/src/component/_generated/api.d.ts

@@ -900,13 +900,13 @@ export type Mounts = {
       "public",
       {
         beforeMessageId?: string;
+        embedding?: Array<number>;
+        embeddingModel?: string;
         limit: number;
         messageRange?: { after: number; before: number };
         searchAllMessagesForUserId?: string;
         text?: string;
         threadId?: string;
-        vector?: Array<number>;
-        vectorModel?: string;
         vectorScoreThreshold?: number;
       },
       Array<{
@@ -1433,6 +1433,12 @@ export type Mounts = {
     >;
   };
   streams: {
+    abort: FunctionReference<
+      "mutation",
+      "public",
+      { reason: string; streamId: string },
+      null
+    >;
     addDelta: FunctionReference<
       "mutation",
       "public",
@@ -1588,13 +1594,18 @@ export type Mounts = {
     list: FunctionReference<
       "query",
       "public",
-      { threadId: string },
+      {
+        startOrder?: number;
+        statuses?: Array<"streaming" | "finished" | "aborted">;
+        threadId: string;
+      },
       Array<{
         agentName?: string;
         model?: string;
         order: number;
         provider?: string;
         providerOptions?: Record<string, Record<string, any>>;
+        status: "streaming" | "finished" | "aborted";
         stepOrder: number;
         streamId: string;
         userId?: string;

package/src/component/messages.ts

@@ -13,7 +13,6 @@ import {
   vMessageStatus,
   vMessageWithMetadataInternal,
   vPaginationResult,
-  vSearchOptions,
 } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import type { Doc, Id } from "./_generated/dataModel.js";
@@ -467,7 +466,14 @@ export const searchMessages = action({
     threadId: v.optional(v.id("threads")),
     searchAllMessagesForUserId: v.optional(v.string()),
     beforeMessageId: v.optional(v.id("messages")),
-    ...vSearchOptions.fields,
+    embedding: v.optional(v.array(v.number())),
+    embeddingModel: v.optional(v.string()),
+    text: v.optional(v.string()),
+    limit: v.number(),
+    vectorScoreThreshold: v.optional(v.number()),
+    messageRange: v.optional(
+      v.object({ before: v.number(), after: v.number() })
+    ),
   },
   returns: v.array(vMessageDoc),
   handler: async (ctx, args): Promise<MessageDoc[]> => {
@@ -486,15 +492,15 @@ export const searchMessages = action({
         beforeMessageId: args.beforeMessageId,
       });
     }
-    if (args.vector) {
-      const dimension = args.vector.length as VectorDimension;
+    if (args.embedding) {
+      const dimension = args.embedding.length as VectorDimension;
       if (!VectorDimensions.includes(dimension)) {
-        throw new Error(`Unsupported vector dimension: ${dimension}`);
+        throw new Error(`Unsupported embedding dimension: ${dimension}`);
       }
       const vectors = (
-        await searchVectors(ctx, args.vector, {
+        await searchVectors(ctx, args.embedding, {
           dimension,
-          model: args.vectorModel ?? "unknown",
+          model: args.embeddingModel ?? "unknown",
           table: "messages",
           searchAllMessagesForUserId: args.searchAllMessagesForUserId,
           threadId: args.threadId,

package/src/component/schema.ts

@@ -113,10 +113,11 @@ export const schema = defineSchema({
       v.object({
         kind: v.literal("finished"),
         endedAt: v.number(),
+        cleanupFnId: v.optional(v.id("_scheduled_functions")),
       }),
       v.object({
-        kind: v.literal("error"),
-        error: v.string(),
+        kind: v.literal("aborted"),
+        reason: v.string(),
       })
     ),
   })

package/src/component/streams.ts

@@ -6,7 +6,7 @@ import {
   vStreamMessage,
 } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
-import type { Id } from "./_generated/dataModel.js";
+import type { Doc, Id } from "./_generated/dataModel.js";
 import {
   internalMutation,
   mutation,
@@ -106,33 +106,87 @@ export const create = mutation({
 export const list = query({
   args: {
     threadId: v.id("threads"),
+    startOrder: v.optional(v.number()),
+    statuses: v.optional(
+      v.array(
+        v.union(
+          v.literal("streaming"),
+          v.literal("finished"),
+          v.literal("aborted")
+        )
+      )
+    ),
   },
   returns: v.array(vStreamMessage),
   handler: async (ctx, args) => {
-    return ctx.db
-      .query("streamingMessages")
-      .withIndex("threadId_state_order_stepOrder", (q) =>
-        q.eq("threadId", args.threadId).eq("state.kind", "streaming")
-      )
-      .order("desc")
-      .take(100)
-      .then((msgs) =>
-        msgs.map((m) => ({
-          streamId: m._id,
-          ...pick(m, [
-            "order",
-            "stepOrder",
-            "userId",
-            "agentName",
-            "model",
-            "provider",
-            "providerOptions",
-          ]),
-        }))
+    const statuses = args.statuses ?? ["streaming"];
+    const messages = await mergedStream(
+      statuses.map((status) =>
+        stream(ctx.db, schema)
+          .query("streamingMessages")
+          .withIndex("threadId_state_order_stepOrder", (q) =>
+            q
+              .eq("threadId", args.threadId)
+              .eq("state.kind", status)
+              .gte("order", args.startOrder ?? 0)
+          )
+          .order("desc")
+      ),
+      ["order", "stepOrder"]
+    ).take(100);
+
+    return messages.map((m) => ({
+      streamId: m._id,
+      status: m.state.kind,
+      ...pick(m, [
+        "order",
+        "stepOrder",
+        "userId",
+        "agentName",
+        "model",
+        "provider",
+        "providerOptions",
+      ]),
+    }));
+  },
+});
+
+export const abort = mutation({
+  args: {
+    streamId: v.id("streamingMessages"),
+    reason: v.string(),
+  },
+  returns: v.null(),
+  handler: async (ctx, args) => {
+    const stream = await ctx.db.get(args.streamId);
+    if (!stream) {
+      throw new Error(`Stream not found: ${args.streamId}`);
+    }
+    if (stream.state.kind !== "streaming") {
+      console.warn(
+        `Stream trying to abort but not currently streaming (${stream.state.kind}): ${args.streamId}`
       );
+      return;
+    }
+    await cleanupTimeoutFn(ctx, stream);
+    await ctx.db.patch(args.streamId, {
+      state: { kind: "aborted", reason: args.reason },
+    });
   },
 });
 
+async function cleanupTimeoutFn(
+  ctx: MutationCtx,
+  stream: Doc<"streamingMessages">
+) {
+  if (stream.state.kind === "streaming" && stream.state.timeoutFnId) {
+    const timeoutFn = await ctx.db.system.get(stream.state.timeoutFnId);
+    if (timeoutFn?.state.kind === "pending") {
+      await ctx.scheduler.cancel(stream.state.timeoutFnId);
+    }
+  }
+}
+
 export const finish = mutation({
   args: {
     streamId: v.id("streamingMessages"),
@@ -153,20 +207,15 @@ export const finish = mutation({
       );
       return;
     }
-    if (stream.state.timeoutFnId) {
-      const timeoutFn = await ctx.db.system.get(stream.state.timeoutFnId);
-      if (timeoutFn?.state.kind === "pending") {
-        await ctx.scheduler.cancel(stream.state.timeoutFnId);
-      }
-    }
-    await ctx.db.patch(args.streamId, {
-      state: { kind: "finished", endedAt: Date.now() },
-    });
-    await ctx.scheduler.runAfter(
+    await cleanupTimeoutFn(ctx, stream);
+    const cleanupFnId = await ctx.scheduler.runAfter(
       DELETE_STREAM_DELAY,
       api.streams.deleteStreamAsync,
       { streamId: args.streamId }
     );
+    await ctx.db.patch(args.streamId, {
+      state: { kind: "finished", endedAt: Date.now(), cleanupFnId },
+    });
   },
 });
 
@@ -223,8 +272,8 @@ export const timeoutStream = internalMutation({
     }
     await ctx.db.patch(args.streamId, {
       state: {
-        kind: "finished",
-        endedAt: Date.now(),
+        kind: "aborted",
+        reason: "timeout",
       },
     });
   },
@@ -245,12 +294,9 @@ async function deletePageForStreamId(
   if (deltas.isDone) {
     const stream = await ctx.db.get(args.streamId);
     if (stream) {
-      const state = stream.state;
-      if (state.kind === "streaming" && state.timeoutFnId) {
-        const timeoutFn = await ctx.db.system.get(state.timeoutFnId);
-        if (timeoutFn?.state.kind === "pending") {
-          await ctx.scheduler.cancel(state.timeoutFnId);
-        }
+      await cleanupTimeoutFn(ctx, stream);
+      if (stream.state.kind === "finished" && stream.state.cleanupFnId) {
+        await ctx.scheduler.cancel(stream.state.cleanupFnId);
       }
       await ctx.db.delete(args.streamId);
     }

package/src/component/vector/index.ts

@@ -182,10 +182,10 @@ export const updateBatch = mutation({
   returns: v.null(),
   handler: async (ctx, args) => {
     await Promise.all(
-      args.vectors.map((v) =>
-        ctx.db.patch(v.id, {
-          model: v.model,
-          vector: v.vector,
+      args.vectors.map((embedding) =>
+        ctx.db.patch(embedding.id, {
+          model: embedding.model,
+          vector: embedding.vector,
         })
       )
     );

package/src/node_modules/.vite/vitest/da39a3ee5e6b4b0d3255bfef95601890afd80709/results.json

@@ -1 +1 @@
-{"version":"3.2.4","results":[[":client/index.test.ts",{"duration":60.55858325958252,"failed":false}],[":component/messages.test.ts",{"duration":27.41837501525879,"failed":false}]]}
+{"version":"3.2.4","results":[[":client/index.test.ts",{"duration":75.65995979309082,"failed":false}],[":component/messages.test.ts",{"duration":43.858792304992676,"failed":false}],[":react/toUIMessages.test.ts",{"duration":4.371167182922363,"failed":false}]]}

package/src/react/deltas.ts

@@ -256,12 +256,14 @@ export function createStreamingMessage(
 ): MessageDoc {
   const { streamId, ...rest } = message;
   const metadata: MessageDoc = {
+    ...rest,
     _id: `${streamId}-${index}`,
     _creationTime: Date.now(),
-    status: "pending",
+    status: (
+      { streaming: "pending", finished: "success", aborted: "failed" } as const
+    )[message.status],
     threadId,
     tool: false,
-    ...rest,
   };
   switch (part.type) {
     case "text-delta":