@convex-dev/agent 0.1.10-alpha.1 → 0.1.11

package/src/client/index.ts

@@ -624,7 +624,7 @@ export class Agent<AgentTools extends ToolSet> {
    * to a thread (and optionally userId).
    */
   async generateObject<T>(
-    ctx: RunActionCtx,
+    ctx: ActionCtx,
     {
       userId: argsUserId,
       threadId,
@@ -705,7 +705,7 @@ export class Agent<AgentTools extends ToolSet> {
    * to a thread (and optionally userId).
    */
   async streamObject<T>(
-    ctx: RunActionCtx,
+    ctx: ActionCtx,
     {
       userId: argsUserId,
       threadId,
@@ -994,8 +994,6 @@ export class Agent<AgentTools extends ToolSet> {
     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;
@@ -1004,8 +1002,6 @@ export class Agent<AgentTools extends ToolSet> {
         kind: "list",
         messages: await ctx.runQuery(this.component.streams.list, {
           threadId: args.threadId,
-          startOrder: args.streamArgs.startOrder,
-          statuses: args.includeStatuses,
         }),
       };
     } else {
@@ -1075,9 +1071,7 @@ export class Agent<AgentTools extends ToolSet> {
       const targetMessage = contextMessages.find(
         (m) => m._id === args.upToAndIncludingMessageId
       )?.message;
-      const messagesToSearch = targetMessage
-        ? [targetMessage, ...args.messages]
-        : args.messages;
+      const messagesToSearch = targetMessage ? [targetMessage] : args.messages;
       if (!("runAction" in ctx)) {
         throw new Error("searchUserMessages only works in an action");
       }
@@ -1351,7 +1345,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @param args The arguments to the saveObject function.
    */
   async saveObject(
-    ctx: RunActionCtx,
+    ctx: ActionCtx,
     args: {
       userId: string | undefined;
       threadId: string;
@@ -1360,10 +1354,15 @@ export class Agent<AgentTools extends ToolSet> {
       metadata?: Omit<MessageWithMetadata, "message">;
     }
   ): Promise<void> {
-    const { messages } = serializeObjectResult(args.result, {
-      model: this.options.chat.modelId,
-      provider: this.options.chat.provider,
-    });
+    const { messages } = await serializeObjectResult(
+      ctx,
+      this.component,
+      args.result,
+      {
+        model: this.options.chat.modelId,
+        provider: this.options.chat.provider,
+      }
+    );
     const embeddings = await this.generateEmbeddings(
       ctx,
       { userId: args.userId, threadId: args.threadId },
@@ -1442,18 +1441,16 @@ export class Agent<AgentTools extends ToolSet> {
   }> {
     contextOptions ||= this.options.contextOptions;
     storageOptions ||= this.options.storageOptions;
-    // If only a messageId is provided, this will be empty.
-    const messages = args.promptMessageId
-      ? []
-      : promptOrMessagesToCoreMessages(args);
+    // If only a promptMessageId is provided, this will be empty.
+    const messages = promptOrMessagesToCoreMessages(args);
     const userId =
       argsUserId ??
       (threadId &&
         (await ctx.runQuery(this.component.threads.getThread, { threadId }))
           ?.userId);
     assert(
-      !args.promptMessageId || !(args.prompt || args.messages),
-      "you can't specify a prompt or message if you specify a promptMessageId"
+      !args.promptMessageId || !args.prompt,
+      "you can't specify a prompt if you specify a promptMessageId"
     );
     // If only a messageId is provided, this will add that message to the end.
     const contextMessages = await this.fetchContextMessages(ctx, {
@@ -1463,29 +1460,33 @@ export class Agent<AgentTools extends ToolSet> {
       messages,
       contextOptions,
     });
-    // Lazily generate embeddings for the prompt message, if it doesn't have
-    // embeddings yet. This can happen if the message was saved in a mutation
-    // where the LLM is not available.
-    if (
-      args.promptMessageId &&
-      !contextMessages.at(-1)?.embeddingId &&
-      this.options.textEmbedding
-    ) {
-      await this.generateAndSaveEmbeddings(ctx, {
-        messageIds: [args.promptMessageId],
-      });
+    // If it was a promptMessageId, pop it off context messages
+    // and add to the end of messages.
+    const promptMessage =
+      !!args.promptMessageId &&
+      contextMessages.at(-1)?._id === args.promptMessageId
+        ? contextMessages.pop()
+        : undefined;
+    if (promptMessage?.message) {
+      messages.push(deserializeMessage(promptMessage.message));
+      // Lazily generate embeddings for the prompt message, if it doesn't have
+      // embeddings yet. This can happen if the message was saved in a mutation
+      // where the LLM is not available.
+      if (!promptMessage.embeddingId && this.options.textEmbedding) {
+        await this.generateAndSaveEmbeddings(ctx, {
+          messageIds: [promptMessage._id],
+        });
+      }
     }
-    let messageId = args.promptMessageId;
-    let order = args.promptMessageId
-      ? contextMessages.at(-1)?.order
-      : undefined;
-    let stepOrder = args.promptMessageId
-      ? contextMessages.at(-1)?.stepOrder
-      : undefined;
+    let messageId = promptMessage?._id;
+    let order = promptMessage?.order;
+    let stepOrder = promptMessage?.stepOrder;
     if (
       threadId &&
       messages.length &&
       storageOptions?.saveMessages !== "none" &&
+      // If it was a promptMessageId, we don't want to save it again.
+      (!args.promptMessageId || storageOptions?.saveMessages === "all") &&
       storageOptions?.saveAnyInputMessages !== false
     ) {
       const saveAll = storageOptions?.saveMessages === "all";
@@ -1494,7 +1495,9 @@ export class Agent<AgentTools extends ToolSet> {
         threadId,
         userId,
         messages: coreMessages,
-        metadata: coreMessages.length === 1 ? [{ id: args.id }] : undefined,
+        metadata: coreMessages.map((_, i) =>
+          i === coreMessages.length - 1 ? { id: args.id } : {}
+        ),
         pending: true,
         failPendingSteps: true,
       });

package/src/client/streaming.ts

@@ -96,13 +96,7 @@ export class DeltaStreamer {
     this.#nextStepOrder = (metadata.stepOrder ?? 0) + 1;
     this.abortController = new AbortController();
     if (metadata.abortSignal) {
-      metadata.abortSignal.addEventListener("abort", async () => {
-        if (this.streamId) {
-          await this.ctx.runMutation(this.component.streams.abort, {
-            streamId: this.streamId,
-            reason: "abortSignal",
-          });
-        }
+      metadata.abortSignal.addEventListener("abort", () => {
         this.abortController.abort();
       });
     }

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

@@ -67,7 +67,7 @@ export type Mounts = {
       "mutation",
       "public",
       { fileIds: Array<string>; force?: boolean },
-      null
+      Array<string>
     >;
     get: FunctionReference<
       "query",
@@ -1414,12 +1414,6 @@ export type Mounts = {
     >;
   };
   streams: {
-    abort: FunctionReference<
-      "mutation",
-      "public",
-      { reason: string; streamId: string },
-      null
-    >;
     addDelta: FunctionReference<
       "mutation",
       "public",
@@ -1575,18 +1569,13 @@ export type Mounts = {
     list: FunctionReference<
       "query",
       "public",
-      {
-        startOrder?: number;
-        statuses?: Array<"streaming" | "finished" | "aborted">;
-        threadId: string;
-      },
+      { 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/files.ts

@@ -149,25 +149,27 @@ export const deleteFiles = mutation({
     fileIds: v.array(v.id("files")),
     force: v.optional(v.boolean()),
   },
+  returns: v.array(v.id("files")),
   handler: async (ctx, args) => {
-    await Promise.all(
+    const deletedFileIds = await Promise.all(
       args.fileIds.map(async (fileId) => {
         const file = await ctx.db.get(fileId);
         if (!file) {
           console.error(`File ${fileId} not found when deleting, skipping...`);
-          return;
+          return null;
         }
         if (file.refcount && file.refcount > 0) {
           if (!args.force) {
             console.error(
               `File ${fileId} has refcount ${file.refcount} > 0, skipping...`
             );
-            return;
+            return null;
           }
         }
         await ctx.db.delete(fileId);
+        return fileId;
       })
     );
+    return deletedFileIds.filter((fileId) => fileId !== null);
   },
-  returns: v.null(),
 });

package/src/component/schema.ts

@@ -108,11 +108,10 @@ 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("aborted"),
-        reason: v.string(),
+        kind: v.literal("error"),
+        error: 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 { Doc, Id } from "./_generated/dataModel.js";
+import type { Id } from "./_generated/dataModel.js";
 import {
   internalMutation,
   mutation,
@@ -106,87 +106,33 @@ 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) => {
-    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 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",
+          ]),
+        }))
       );
-      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"),
@@ -207,15 +153,20 @@ export const finish = mutation({
       );
       return;
     }
-    await cleanupTimeoutFn(ctx, stream);
-    const cleanupFnId = await ctx.scheduler.runAfter(
+    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(
       DELETE_STREAM_DELAY,
       api.streams.deleteStreamAsync,
       { streamId: args.streamId }
     );
-    await ctx.db.patch(args.streamId, {
-      state: { kind: "finished", endedAt: Date.now(), cleanupFnId },
-    });
   },
 });
 
@@ -272,8 +223,8 @@ export const timeoutStream = internalMutation({
     }
     await ctx.db.patch(args.streamId, {
       state: {
-        kind: "aborted",
-        reason: "timeout",
+        kind: "finished",
+        endedAt: Date.now(),
       },
     });
   },
@@ -294,9 +245,12 @@ async function deletePageForStreamId(
   if (deltas.isDone) {
     const stream = await ctx.db.get(args.streamId);
     if (stream) {
-      await cleanupTimeoutFn(ctx, stream);
-      if (stream.state.kind === "finished" && stream.state.cleanupFnId) {
-        await ctx.scheduler.cancel(stream.state.cleanupFnId);
+      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 ctx.db.delete(args.streamId);
     }

package/src/component/vector/tables.ts

@@ -15,7 +15,7 @@ import {
   type VObject,
   type VUnion,
 } from "convex/values";
-import type { QueryCtx } from "../_generated/server";
+import type { QueryCtx } from "../_generated/server.js";
 
 // We only generate embeddings for non-tool, non-system messages
 const embeddings = {

package/src/mapping.ts

@@ -85,6 +85,7 @@ export async function serializeNewMessagesInStep<TOOLS extends ToolSet>(
     usage: step.usage,
     warnings: step.warnings,
     finishReason: step.finishReason,
+    sources: step.stepType === "tool-result" ? undefined : step.sources,
   };
   const toolFields = {
     sources: step.sources,
@@ -113,17 +114,19 @@ export async function serializeNewMessagesInStep<TOOLS extends ToolSet>(
   return messages;
 }
 
-export function serializeObjectResult(
+export async function serializeObjectResult(
+  ctx: ActionCtx,
+  component: AgentComponent,
   result: GenerateObjectResult<unknown>,
   metadata: { model: string; provider: string }
-): { messages: MessageWithMetadata[] } {
+): Promise<{ messages: MessageWithMetadata[] }> {
   const text = JSON.stringify(result.object);
 
-  const message = {
+  const { message, fileIds } = await serializeMessage(ctx, component, {
     role: "assistant" as const,
     content: text,
     id: result.response.id,
-  };
+  });
   return {
     messages: [
       {
@@ -136,6 +139,7 @@ export function serializeObjectResult(
         text,
         usage: result.usage,
         warnings: result.warnings,
+        fileIds,
       },
     ],
   };
@@ -330,9 +334,17 @@ function encodeBase64(data: ArrayBuffer): string {
 export function promptOrMessagesToCoreMessages(args: {
   prompt?: string;
   messages?: CoreMessage[] | AIMessageWithoutId[];
+  promptMessageId?: string;
 }): CoreMessage[] {
   const messages: CoreMessage[] = [];
-  assert(args.prompt || args.messages, "messages or prompt is required");
+  assert(
+    args.prompt || args.messages || args.promptMessageId,
+    "messages or prompt or promptMessageId is required"
+  );
+  assert(
+    !args.promptMessageId || !args.prompt,
+    "you can't specify a prompt if you specify a promptMessageId. use messages instead."
+  );
   if (args.messages) {
     if (
       args.messages.some(
@@ -353,6 +365,5 @@ export function promptOrMessagesToCoreMessages(args: {
   if (args.prompt) {
     messages.push({ role: "user", content: args.prompt });
   }
-  assert(messages.length > 0, "Messages must contain at least one message");
   return messages;
 }

package/src/react/deltas.ts

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

package/src/react/index.ts

@@ -1,5 +1,5 @@
 "use client";
-import { omit, type ErrorMessage } from "convex-helpers";
+import type { ErrorMessage } from "convex-helpers";
 import {
   type PaginatedQueryArgs,
   type UsePaginatedQueryResult,
@@ -7,7 +7,7 @@ import {
 } from "convex/react";
 import { usePaginatedQuery } from "convex-helpers/react";
 import type { FunctionArgs } from "convex/server";
-import { useMemo, useRef, useState } from "react";
+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";
@@ -103,18 +103,14 @@ export function useThreadMessages<
       ThreadMessagesArgs<Query>,
       ThreadMessagesResult<Query>
     >,
-    !options.stream ||
-      args === "skip" ||
-      paginated.status === "LoadingFirstPage"
-      ? "skip"
-      : { ...args, startOrder: paginated.results.at(-1)?.order }
+    !options.stream ? "skip" : args
   );
 
   const merged = useMemo(() => {
     const streamListMessages =
       streamMessages?.map((m) => ({
         ...m,
-        streaming: !m.status || m.status === "streaming",
+        streaming: true,
       })) ?? [];
     return {
       ...paginated,
@@ -154,7 +150,7 @@ export function useStreamingThreadMessages<
   Query extends ThreadStreamQuery<any, any>,
 >(
   query: Query,
-  args: (ThreadMessagesArgs<Query> & { startOrder?: number }) | "skip"
+  args: ThreadMessagesArgs<Query> | "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
@@ -162,23 +158,15 @@ export function useStreamingThreadMessages<
   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
+    args === "skip"
+      ? args
       : ({
-          ...queryArgs,
+          ...args,
           paginationOpts: { cursor: null, numItems: 0 },
-          streamArgs: {
-            kind: "list",
-            startOrder: startOrderRef.current,
-          } as StreamArgs,
+          streamArgs: { kind: "list" } as StreamArgs,
         } as FunctionArgs<Query>)
   ) as
     | { streams: Extract<SyncStreamsReturnValue, { kind: "list" }> }
@@ -198,10 +186,10 @@ export function useStreamingThreadMessages<
   // Get the deltas for all the active streams, if any.
   const cursorQuery = useQuery(
     query,
-    queryArgs === "skip" || !streamList
+    args === "skip" || !streamList
       ? ("skip" as const)
       : ({
-          ...queryArgs,
+          ...args,
           paginationOpts: { cursor: null, numItems: 0 },
           streamArgs: { kind: "deltas", cursors } as StreamArgs,
         } as FunctionArgs<Query>)

package/src/react/optimisticallySendMessage.ts

@@ -11,7 +11,7 @@ export function optimisticallySendMessage(
 ) => void {
   return (store, args) => {
     const queries = store.getAllQueries(query);
-    let maxOrder = -1;
+    let maxOrder = 0;
     let maxStepOrder = 0;
     for (const q of queries) {
       if (q.args?.threadId !== args.threadId) continue;

package/src/validators.ts

@@ -460,7 +460,6 @@ export const vStreamArgs = v.optional(
   v.union(
     v.object({
       kind: v.literal("list"),
-      startOrder: v.optional(v.number()),
     }),
     v.object({
       kind: v.literal("deltas"),
@@ -472,11 +471,6 @@ export type StreamArgs = Infer<typeof vStreamArgs>;
 
 export const vStreamMessage = v.object({
   streamId: v.string(),
-  status: v.union(
-    v.literal("streaming"),
-    v.literal("finished"),
-    v.literal("aborted")
-  ),
   order: v.number(),
   stepOrder: v.number(),
   // metadata

package/dist/commonjs/client/_generated/_ignore.d.ts

@@ -1 +0,0 @@
-//# sourceMappingURL=_ignore.d.ts.map

package/dist/commonjs/client/_generated/_ignore.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"_ignore.d.ts","sourceRoot":"","sources":["../../../../src/client/_generated/_ignore.ts"],"names":[],"mappings":""}

package/dist/commonjs/client/_generated/_ignore.js

@@ -1,3 +0,0 @@
-"use strict";
-// This is only here so convex-test can detect a _generated folder
-//# sourceMappingURL=_ignore.js.map

package/dist/commonjs/client/_generated/_ignore.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"_ignore.js","sourceRoot":"","sources":["../../../../src/client/_generated/_ignore.ts"],"names":[],"mappings":";AAAA,kEAAkE"}