@convex-dev/agent 0.1.9-alpha.1 → 0.1.9-alpha.11

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

@@ -1414,6 +1414,12 @@ export type Mounts = {
     >;
   };
   streams: {
+    abort: FunctionReference<
+      "mutation",
+      "public",
+      { reason: string; streamId: string },
+      null
+    >;
     addDelta: FunctionReference<
       "mutation",
       "public",
@@ -1467,7 +1473,7 @@ export type Mounts = {
         start: number;
         streamId: string;
       },
-      null
+      boolean
     >;
     create: FunctionReference<
       "mutation",
@@ -1569,13 +1575,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;
@@ -1663,13 +1674,21 @@ export type Mounts = {
     deleteAllForThreadIdAsync: FunctionReference<
       "mutation",
       "public",
-      { cursor?: string; limit?: number; threadId: string },
-      { cursor: string; isDone: boolean }
+      {
+        cursor?: string;
+        deltaCursor?: string;
+        limit?: number;
+        messagesDone?: boolean;
+        streamOrder?: number;
+        streamsDone?: boolean;
+        threadId: string;
+      },
+      { isDone: boolean }
     >;
     deleteAllForThreadIdSync: FunctionReference<
       "action",
       "public",
-      { cursor?: string; limit?: number; threadId: string },
+      { limit?: number; threadId: string },
       null
     >;
     getThread: FunctionReference<

package/src/component/schema.ts

@@ -110,8 +110,8 @@ export const schema = defineSchema({
         endedAt: v.number(),
       }),
       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,
@@ -17,20 +17,30 @@ import {
 import schema from "./schema.js";
 import { stream } from "convex-helpers/server/stream";
 import { mergedStream } from "convex-helpers/server/stream";
+import { paginator } from "convex-helpers/server/pagination";
+
+const SECOND = 1000;
+const MINUTE = 60 * SECOND;
 
 const MAX_DELTAS_PER_REQUEST = 1000;
 const MAX_DELTAS_PER_STREAM = 100;
-const TIMEOUT_INTERVAL = 1000 * 60; // 1 minute
-const DELETE_STREAM_DELAY = 1000 * 60 * 5; // 5 minutes
+const TIMEOUT_INTERVAL = 10 * MINUTE;
+const DELETE_STREAM_DELAY = MINUTE * 5; // 5 minutes
 
 const deltaValidator = schema.tables.streamDeltas.validator;
 
 export const addDelta = mutation({
   args: deltaValidator,
-  returns: v.null(),
+  returns: v.boolean(),
   handler: async (ctx, args) => {
     await ctx.db.insert("streamDeltas", args);
     await heartbeatStream(ctx, { streamId: args.streamId });
+    const stream = await ctx.db.get(args.streamId);
+    if (stream?.state.kind !== "streaming") {
+      console.warn(`Stream is not streaming: ${args.streamId}`);
+      return false;
+    }
+    return true;
   },
 });
 
@@ -96,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"),
@@ -138,14 +202,12 @@ export const finish = mutation({
       throw new Error(`Stream not found: ${args.streamId}`);
     }
     if (stream.state.kind !== "streaming") {
-      throw new Error(`Stream is not streaming: ${args.streamId}`);
-    }
-    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);
-      }
+      console.warn(
+        `Stream trying to finish but not currently streaming: ${args.streamId}`
+      );
+      return;
     }
+    await cleanupTimeoutFn(ctx, stream);
     await ctx.db.patch(args.streamId, {
       state: { kind: "finished", endedAt: Date.now() },
     });
@@ -210,8 +272,8 @@ export const timeoutStream = internalMutation({
     }
     await ctx.db.patch(args.streamId, {
       state: {
-        kind: "finished",
-        endedAt: Date.now(),
+        kind: "aborted",
+        reason: "timeout",
       },
     });
   },
@@ -221,7 +283,7 @@ async function deletePageForStreamId(
   ctx: MutationCtx,
   args: { streamId: Id<"streamingMessages">; cursor?: string }
 ) {
-  const deltas = await ctx.db
+  const deltas = await paginator(ctx.db, schema)
     .query("streamDeltas")
     .withIndex("streamId_start_end", (q) => q.eq("streamId", args.streamId))
     .paginate({
@@ -229,6 +291,9 @@ async function deletePageForStreamId(
       cursor: args.cursor ?? null,
     });
   await Promise.all(deltas.page.map((d) => ctx.db.delete(d._id)));
+  if (deltas.isDone) {
+    await ctx.db.delete(args.streamId);
+  }
   return deltas;
 }
 
@@ -268,13 +333,12 @@ export async function deleteStreamsPageForThreadId(
     cursor: deltaCursor,
   });
   if (result.isDone) {
-    await ctx.db.delete(streamMessage._id);
     deltaCursor = undefined;
   }
   return {
     isDone: false,
     streamOrder: streamMessage.order,
-    deltaCursor: result.continueCursor,
+    deltaCursor,
   };
 }
 
@@ -333,7 +397,6 @@ export const deleteStreamSync = mutation({
         cursor: deltas.continueCursor,
       });
     }
-    await ctx.db.delete(args.streamId);
   },
 });
 

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/deltas.ts

@@ -189,6 +189,7 @@ export function applyDeltasToStreamMessage(
         contentToAdd = part;
         break;
       case "reasoning":
+        currentMessage.reasoning += part.textDelta;
         if (lastContent?.type === "reasoning") {
           lastContent.text += part.textDelta;
         } else {
@@ -255,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":

package/src/react/index.ts

@@ -1,5 +1,5 @@
 "use client";
-import type { ErrorMessage } from "convex-helpers";
+import { omit, 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, 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,14 +103,18 @@ export function useThreadMessages<
       ThreadMessagesArgs<Query>,
       ThreadMessagesResult<Query>
     >,
-    !options.stream ? "skip" : args
+    !options.stream ||
+      args === "skip" ||
+      paginated.status === "LoadingFirstPage"
+      ? "skip"
+      : { ...args, startOrder: paginated.results.at(-1)?.order }
   );
 
   const merged = useMemo(() => {
     const streamListMessages =
       streamMessages?.map((m) => ({
         ...m,
-        streaming: true,
+        streaming: !m.status || m.status === "streaming",
       })) ?? [];
     return {
       ...paginated,
@@ -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,15 +162,23 @@ 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,
-    args === "skip"
-      ? args
+    queryArgs === "skip"
+      ? queryArgs
       : ({
-          ...args,
+          ...queryArgs,
           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" }> }
@@ -186,10 +198,10 @@ export function useStreamingThreadMessages<
   // Get the deltas for all the active streams, if any.
   const cursorQuery = useQuery(
     query,
-    args === "skip" || !streamList
+    queryArgs === "skip" || !streamList
       ? ("skip" as const)
       : ({
-          ...args,
+          ...queryArgs,
           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 = 0;
+    let maxOrder = -1;
     let maxStepOrder = 0;
     for (const q of queries) {
       if (q.args?.threadId !== args.threadId) continue;

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"),
@@ -471,6 +472,11 @@ 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