@convex-dev/agent 0.1.9-alpha.2 → 0.1.9-alpha.4

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,
@@ -102,13 +102,17 @@ export const create = mutation({
 export const list = query({
   args: {
     threadId: v.id("threads"),
+    startOrder: v.optional(v.number()),
   },
   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")
+        q
+          .eq("threadId", args.threadId)
+          .eq("state.kind", "streaming")
+          .gte("order", args.startOrder ?? 0)
       )
       .order("desc")
       .take(100)
@@ -129,6 +133,42 @@ export const list = query({
   },
 });
 
+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"),
@@ -149,12 +189,7 @@ 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 cleanupTimeoutFn(ctx, stream);
     await ctx.db.patch(args.streamId, {
       state: { kind: "finished", endedAt: Date.now() },
     });

package/src/component/threads.ts

@@ -16,6 +16,7 @@ import {
 } from "./_generated/server.js";
 import { deleteMessage } from "./messages.js";
 import { schema, v } from "./schema.js";
+import { deleteStreamsPageForThreadId } from "./streams.js";
 
 function publicThreadOrNull(thread: Doc<"threads"> | null): ThreadDoc | null {
   if (thread === null) {
@@ -100,29 +101,17 @@ export const updateThread = mutation({
 
 // TODO: delete thread
 
-const deleteThreadArgs = {
-  threadId: v.id("threads"),
-  cursor: v.optional(v.string()),
-  limit: v.optional(v.number()),
-};
-type DeleteThreadArgs = ObjectType<typeof deleteThreadArgs>;
-const deleteThreadReturns = {
-  cursor: v.string(),
-  isDone: v.boolean(),
-};
-type DeleteThreadReturns = ObjectType<typeof deleteThreadReturns>;
-
 /**
  * Use this to delete a thread and everything it contains.
  * It will try to delete all pages synchronously.
  * If it times out or fails, you'll have to run it again.
  */
 export const deleteAllForThreadIdSync = action({
-  args: deleteThreadArgs,
+  args: { threadId: v.id("threads"), limit: v.optional(v.number()) },
   handler: async (ctx, args) => {
-    let cursor = args.cursor;
+    let cursor: string | undefined = undefined;
     while (true) {
-      const result = await ctx.runMutation(
+      const result: DeleteThreadReturns = await ctx.runMutation(
         internal.threads._deletePageForThreadId,
         { threadId: args.threadId, cursor, limit: args.limit }
       );
@@ -138,6 +127,22 @@ export const deleteAllForThreadIdSync = action({
   returns: v.null(),
 });
 
+const deleteThreadArgs = {
+  threadId: v.id("threads"),
+  cursor: v.optional(v.string()),
+  messagesDone: v.optional(v.boolean()),
+  streamsDone: v.optional(v.boolean()),
+  streamOrder: v.optional(v.number()),
+  deltaCursor: v.optional(v.string()),
+  limit: v.optional(v.number()),
+};
+type DeleteThreadArgs = ObjectType<typeof deleteThreadArgs>;
+const deleteThreadReturns = {
+  cursor: v.string(),
+  isDone: v.boolean(),
+};
+type DeleteThreadReturns = ObjectType<typeof deleteThreadReturns>;
+
 export const _deletePageForThreadId = internalMutation({
   args: deleteThreadArgs,
   handler: deletePageForThreadIdHandler,
@@ -151,22 +156,39 @@ export const _deletePageForThreadId = internalMutation({
 export const deleteAllForThreadIdAsync = mutation({
   args: deleteThreadArgs,
   handler: async (ctx, args) => {
-    const result = await deletePageForThreadIdHandler(ctx, args);
-    if (!result.isDone) {
+    let messagesResult = {
+      isDone: args.messagesDone ?? false,
+      cursor: args.cursor,
+    };
+    if (!args.messagesDone) {
+      messagesResult = await deletePageForThreadIdHandler(ctx, args);
+    }
+    let streamResult = {
+      isDone: args.streamsDone ?? false,
+      streamOrder: args.streamOrder,
+      deltaCursor: args.deltaCursor,
+    };
+    if (!args.streamsDone) {
+      streamResult = await deleteStreamsPageForThreadId(ctx, {
+        threadId: args.threadId,
+        streamOrder: args.streamOrder,
+        deltaCursor: args.deltaCursor,
+      });
+    }
+    const isDone = messagesResult.isDone && streamResult.isDone;
+    if (!isDone) {
       await ctx.scheduler.runAfter(0, api.threads.deleteAllForThreadIdAsync, {
         threadId: args.threadId,
-        cursor: result.cursor,
+        cursor: messagesResult.cursor,
+        messagesDone: messagesResult.isDone,
+        streamsDone: streamResult.isDone,
+        streamOrder: streamResult.streamOrder,
+        deltaCursor: streamResult.deltaCursor,
       });
     }
-    // Kick off the streams deletion
-    await ctx.scheduler.runAfter(
-      0,
-      api.streams.deleteAllStreamsForThreadIdSync,
-      { threadId: args.threadId }
-    );
-    return result;
+    return { isDone };
   },
-  returns: deleteThreadReturns,
+  returns: v.object({ isDone: v.boolean() }),
 });
 
 async function deletePageForThreadIdHandler(

package/src/react/index.ts

@@ -7,7 +7,7 @@ import {
 } from "convex/react";
 import { usePaginatedQuery } from "convex-helpers/react";
 import type { FunctionArgs } from "convex/server";
-import { useMemo, useState } from "react";
+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";
@@ -103,7 +103,11 @@ export function useThreadMessages<
       ThreadMessagesArgs<Query>,
       ThreadMessagesResult<Query>
     >,
-    !options.stream ? "skip" : args
+    !options.stream ||
+      args === "skip" ||
+      paginated.status === "LoadingFirstPage"
+      ? "skip"
+      : { ...args, startOrder: paginated.results.at(0)?.order }
   );
 
   const merged = useMemo(() => {
@@ -150,7 +154,7 @@ export function useStreamingThreadMessages<
   Query extends ThreadStreamQuery<any, any>,
 >(
   query: Query,
-  args: ThreadMessagesArgs<Query> | "skip"
+  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
@@ -158,6 +162,10 @@ export function useStreamingThreadMessages<
   const [streams, setStreams] = useState<
     Array<{ streamId: string; cursor: number; messages: MessageDoc[] }>
   >([]);
+  const startOrderRef = useRef<number>(0);
+  if (args !== "skip" && !startOrderRef.current && args.startOrder) {
+    startOrderRef.current = args.startOrder;
+  }
   // Get all the active streams
   const streamList = useQuery(
     query,
@@ -166,7 +174,10 @@ export function useStreamingThreadMessages<
       : ({
           ...args,
           paginationOpts: { cursor: null, numItems: 0 },
-          streamArgs: { kind: "list" } as StreamArgs,
+          streamArgs: {
+            kind: "list",
+            startOrder: startOrderRef.current,
+          } as StreamArgs,
         } as FunctionArgs<Query>)
   ) as
     | { streams: Extract<SyncStreamsReturnValue, { kind: "list" }> }

package/src/validators.ts

@@ -460,6 +460,7 @@ export const vStreamArgs = v.optional(
   v.union(
     v.object({
       kind: v.literal("list"),
+      startOrder: v.optional(v.number()),
     }),
     v.object({
       kind: v.literal("deltas"),