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

package/src/client/streaming.ts

@@ -1,26 +1,40 @@
 import {
-  type ChunkDetector,
   smoothStream,
+  type AsyncIterableStream,
+  type ChunkDetector,
   type StreamTextTransform,
+  type TextStreamPart,
   type ToolSet,
+  type UIMessageChunk,
 } from "ai";
-import type {
-  ProviderOptions,
-  StreamArgs,
-  StreamDelta,
-  StreamMessage,
-  vTextStreamPartV5,
+import { v } from "convex/values";
+import {
+  vMessageDoc,
+  vPaginationResult,
+  vStreamDelta,
+  vStreamMessage,
+  type ProviderOptions,
+  type StreamArgs,
+  type StreamDelta,
+  type StreamMessage,
 } from "../validators.js";
-import type { MessageDoc } from "../component/schema.js";
 import type {
+  ActionCtx,
   AgentComponent,
-  RunActionCtx,
-  RunMutationCtx,
-  RunQueryCtx,
+  MutationCtx,
+  QueryCtx,
   SyncStreamsReturnValue,
 } from "./types.js";
-import { omit } from "convex-helpers";
-import type { Infer } from "convex/values";
+
+export const vStreamMessagesReturnValue = v.object({
+  ...vPaginationResult(vMessageDoc).fields,
+  streams: v.optional(
+    v.union(
+      v.object({ kind: v.literal("list"), messages: v.array(vStreamMessage) }),
+      v.object({ kind: v.literal("deltas"), deltas: v.array(vStreamDelta) }),
+    ),
+  ),
+});
 
 /**
  * A function that handles fetching stream deltas, used with the React hooks
@@ -32,42 +46,47 @@ import type { Infer } from "convex/values";
  * @returns The deltas for each stream from their existing cursor.
  */
 export async function syncStreams(
-  ctx: RunQueryCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
-  args: {
+  {
+    threadId,
+    streamArgs,
+    includeStatuses,
+  }: {
     threadId: string;
-    streamArgs: StreamArgs | undefined;
+    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") {
+  if (!streamArgs) return undefined;
+  if (streamArgs.kind === "list") {
     return {
       kind: "list",
       messages: await listStreams(ctx, component, {
-        threadId: args.threadId,
-        startOrder: args.streamArgs.startOrder,
-        includeStatuses: args.includeStatuses,
+        threadId,
+        startOrder: streamArgs.startOrder,
+        includeStatuses,
       }),
     };
   } else {
     return {
       kind: "deltas",
       deltas: await ctx.runQuery(component.streams.listDeltas, {
-        threadId: args.threadId,
-        cursors: args.streamArgs.cursors,
+        threadId,
+        cursors: streamArgs.cursors,
       }),
     };
   }
 }
 
 export async function abortStream(
-  ctx: RunMutationCtx,
+  ctx: MutationCtx | ActionCtx,
   component: AgentComponent,
-  args: {
-    reason: string;
-  } & ({ streamId: string } | { threadId: string; order: number }),
+  args: { reason: string } & (
+    | { streamId: string }
+    | { threadId: string; order: number }
+  ),
 ): Promise<boolean> {
   if ("streamId" in args) {
     return await ctx.runMutation(component.streams.abort, {
@@ -93,7 +112,7 @@ export async function abortStream(
  * @returns The streams for the thread.
  */
 export async function listStreams(
-  ctx: RunQueryCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
   {
     threadId,
@@ -126,15 +145,31 @@ export type StreamingOptions = {
    * Defaults to 250.
    */
   throttleMs?: number;
+  /**
+   * If set to true, this will return immediately, as it would if you weren't
+   * saving the deltas. Otherwise, the call will "consume" the stream with
+   * .consumeStream(), which waits for the stream to finish before returning.
+   *
+   * When saving deltas, you're often not interactin with the stream otherwise.
+   */
+  returnImmediately?: boolean;
 };
 export const DEFAULT_STREAMING_OPTIONS = {
   // This chunks by sentences / clauses. Punctuation followed by whitespace.
   chunking: /[\p{P}\s]/u,
   throttleMs: 250,
+  returnImmediately: false,
 } satisfies StreamingOptions;
 
+/**
+ *
+ * @param options The options passed to `agent.streamText` to decide whether to
+ * save deltas while streaming.
+ * @param existing The transforms passed to `agent.streamText` to merge with.
+ * @returns The merged transforms to pass to the underlying `streamText` call.
+ */
 export function mergeTransforms<TOOLS extends ToolSet>(
-  options: StreamingOptions | boolean | undefined,
+  options: { chunking?: StreamingOptions["chunking"] } | boolean | undefined,
   existing:
     | StreamTextTransform<TOOLS>
     | Array<StreamTextTransform<TOOLS>>
@@ -156,84 +191,138 @@ export function mergeTransforms<TOOLS extends ToolSet>(
   return transforms;
 }
 
-export class DeltaStreamer {
-  public streamId: string | undefined;
-  public readonly options: Required<StreamingOptions>;
-  #nextParts: Infer<typeof vTextStreamPartV5>[] = [];
-  #nextOrder: number;
-  #nextStepOrder: number;
+/**
+ * DeltaStreamer can be used to save a stream of "parts" by writing
+ * batches of them in "deltas" to the database so clients can subscribe
+ * (using the syncStreams utility and client hooks) and re-hydrate the stream.
+ * You can optionally compress the parts, e.g. concatenating text deltas, to
+ * optimize the data in transit.
+ */
+export class DeltaStreamer<T> {
+  streamId: string | undefined;
+  public readonly config: {
+    throttleMs: number;
+    onAsyncAbort: (reason: string) => Promise<void>;
+    compress: ((parts: T[]) => T[]) | null;
+  };
+  #nextParts: T[] = [];
   #latestWrite: number = 0;
   #ongoingWrite: Promise<void> | undefined;
   #cursor: number = 0;
   public abortController: AbortController;
+  // When true, the stream will be finished externally (e.g., atomically via addMessages)
+  // and consumeStream should skip calling finish().
+  #finishedExternally: boolean = false;
 
   constructor(
     public readonly component: AgentComponent,
-    public readonly ctx: RunActionCtx,
-    options: true | StreamingOptions,
+    public readonly ctx: MutationCtx | ActionCtx,
+    config: {
+      throttleMs: number | undefined;
+      onAsyncAbort: (reason: string) => Promise<void>;
+      abortSignal: AbortSignal | undefined;
+      compress: ((parts: T[]) => T[]) | null;
+    },
     public readonly metadata: {
       threadId: string;
-      agentName: string | undefined;
-      model: string | undefined;
-      provider: string | undefined;
-      providerOptions: ProviderOptions | undefined;
-      userId: string | undefined;
-      order: number | undefined;
-      stepOrder: number | undefined;
-      abortSignal: AbortSignal | undefined;
+      userId?: string;
+      order: number;
+      stepOrder: number;
+      agentName?: string;
+      model?: string;
+      provider?: string;
+      providerOptions?: ProviderOptions;
+      format: "UIMessageChunk" | "TextStreamPart" | undefined;
     },
   ) {
-    this.options =
-      typeof options === "boolean"
-        ? DEFAULT_STREAMING_OPTIONS
-        : {
-            ...DEFAULT_STREAMING_OPTIONS,
-            ...options,
-          };
+    this.config = {
+      throttleMs: config.throttleMs ?? DEFAULT_STREAMING_OPTIONS.throttleMs,
+      onAsyncAbort: config.onAsyncAbort,
+      compress: config.compress,
+    };
     this.#nextParts = [];
-    this.#nextOrder = metadata.order ?? 0;
-    this.#nextStepOrder = (metadata.stepOrder ?? 0) + 1;
     this.abortController = new AbortController();
-    if (metadata.abortSignal) {
-      metadata.abortSignal.addEventListener("abort", async () => {
+    if (config.abortSignal) {
+      config.abortSignal.addEventListener("abort", async () => {
+        if (this.abortController.signal.aborted) {
+          return;
+        }
         if (this.streamId) {
+          this.abortController.abort();
+          await this.#ongoingWrite;
           await this.ctx.runMutation(this.component.streams.abort, {
             streamId: this.streamId,
             reason: "abortSignal",
           });
         }
-        this.abortController.abort();
       });
     }
   }
-  public async addParts(parts: Infer<typeof vTextStreamPartV5>[]) {
+
+  // Avoid race conditions by only creating once
+  #creatingStreamIdPromise: Promise<string> | undefined;
+  public async getStreamId() {
+    if (!this.streamId) {
+      if (!this.#creatingStreamIdPromise) {
+        this.#creatingStreamIdPromise = this.ctx.runMutation(
+          this.component.streams.create,
+          this.metadata,
+        );
+      }
+      this.streamId = await this.#creatingStreamIdPromise;
+    }
+    return this.streamId;
+  }
+
+  public async addParts(parts: T[]) {
     if (this.abortController.signal.aborted) {
       return;
     }
-    if (!this.streamId) {
-      this.streamId = await this.ctx.runMutation(
-        this.component.streams.create,
-        {
-          ...omit(this.metadata, ["abortSignal"]),
-          order: this.#nextOrder,
-          stepOrder: this.#nextStepOrder,
-        },
-      );
-    }
+    await this.getStreamId();
     this.#nextParts.push(...parts);
     if (
       !this.#ongoingWrite &&
-      Date.now() - this.#latestWrite >= this.options.throttleMs
+      Date.now() - this.#latestWrite >= this.config.throttleMs
     ) {
       this.#ongoingWrite = this.#sendDelta();
     }
   }
 
+  public async consumeStream(stream: AsyncIterableStream<T>) {
+    for await (const chunk of stream) {
+      await this.addParts([chunk]);
+    }
+    // Skip finish if it will be handled externally (atomically with message save)
+    if (!this.#finishedExternally) {
+      await this.finish();
+    }
+  }
+
+  /**
+   * Mark the stream as being finished externally (e.g., atomically via addMessages).
+   * When called, consumeStream() will skip calling finish() since it will be
+   * handled elsewhere in the same mutation as message saving.
+   */
+  public markFinishedExternally(): void {
+    this.#finishedExternally = true;
+  }
+
+  /**
+   * Get the stream ID, waiting for it to be created if necessary.
+   * Useful for passing to addMessages for atomic finish.
+   */
+  public async getOrCreateStreamId(): Promise<string> {
+    return this.getStreamId();
+  }
+
   async #sendDelta() {
     if (this.abortController.signal.aborted) {
       return;
     }
     const delta = this.#createDelta();
+    if (!delta) {
+      return;
+    }
     this.#latestWrite = Date.now();
     try {
       const success = await this.ctx.runMutation(
@@ -241,16 +330,21 @@ export class DeltaStreamer {
         delta,
       );
       if (!success) {
+        await this.config.onAsyncAbort("async abort");
         this.abortController.abort();
+        return;
       }
     } catch (e) {
+      await this.config.onAsyncAbort(
+        e instanceof Error ? e.message : "unknown error",
+      );
       this.abortController.abort();
       throw e;
     }
     // Now that we've sent the delta, check if we need to send another one.
     if (
       this.#nextParts.length > 0 &&
-      Date.now() - this.#latestWrite >= this.options.throttleMs
+      Date.now() - this.#latestWrite >= this.config.throttleMs
     ) {
       // We send again immediately with the accumulated deltas.
       this.#ongoingWrite = this.#sendDelta();
@@ -259,47 +353,97 @@ export class DeltaStreamer {
     }
   }
 
-  #createDelta(): StreamDelta {
+  #createDelta(): StreamDelta | undefined {
+    if (this.#nextParts.length === 0) {
+      return undefined;
+    }
     const start = this.#cursor;
     const end = start + this.#nextParts.length;
     this.#cursor = end;
-    const parts = this.#nextParts;
+    const parts = this.config.compress
+      ? this.config.compress(this.#nextParts)
+      : this.#nextParts;
     this.#nextParts = [];
     if (!this.streamId) {
       throw new Error("Creating a delta before the stream is created");
     }
-    return {
+    return { streamId: this.streamId, start, end, parts };
+  }
+
+  public async finish() {
+    if (!this.streamId) {
+      return;
+    }
+    await this.#ongoingWrite;
+    await this.#sendDelta();
+    await this.ctx.runMutation(this.component.streams.finish, {
       streamId: this.streamId,
-      start,
-      end,
-      parts,
-    };
+    });
   }
 
-  public async finish(messages: MessageDoc[]) {
-    if (this.#ongoingWrite) {
-      await this.#ongoingWrite;
-      this.#ongoingWrite = undefined;
+  public async fail(reason: string) {
+    if (this.abortController.signal.aborted) {
+      return;
     }
+    this.abortController.abort();
     if (!this.streamId) {
-      throw new Error("Finish called before stream is created");
+      return;
     }
-    const lastMessage = messages.at(-1);
-    if (lastMessage) {
-      this.#nextOrder = lastMessage.order;
-      this.#nextStepOrder = lastMessage.stepOrder + 1;
+    await this.#ongoingWrite;
+    await this.ctx.runMutation(this.component.streams.abort, {
+      streamId: this.streamId,
+      reason,
+    });
+  }
+}
+
+/**
+ * Compressing parts when streaming to save bandwidth in deltas.
+ */
+
+export function compressUIMessageChunks(
+  parts: UIMessageChunk[],
+): UIMessageChunk[] {
+  const compressed: UIMessageChunk[] = [];
+  for (const part of parts) {
+    const last = compressed.at(-1);
+    if (part.type === "text-delta" || part.type === "reasoning-delta") {
+      if (last?.type === part.type && part.id === last.id) {
+        last.delta += part.delta;
+      } else {
+        compressed.push(part);
+      }
     } else {
-      console.warn("Step finished without generating a message");
+      compressed.push(part);
+    }
+  }
+  return compressed;
+}
+
+export function compressTextStreamParts(
+  parts: TextStreamPart<ToolSet>[],
+): TextStreamPart<ToolSet>[] {
+  const compressed: TextStreamPart<ToolSet>[] = [];
+  for (const part of parts) {
+    const last = compressed.at(-1);
+    if (part.type === "text-delta" || part.type === "reasoning-delta") {
+      if (last?.type === part.type && part.id === last.id) {
+        last.text += part.text;
+      } else {
+        compressed.push(part);
+      }
+    } else {
+      if (part.type === "file") {
+        compressed.push({
+          type: "file",
+          file: {
+            ...part.file,
+            uint8Array: undefined as unknown as Uint8Array,
+          },
+        });
+      }
+      compressed.push(part);
     }
-    const finalDelta =
-      this.#nextParts.length > 0 ? this.#createDelta() : undefined;
-    this.#nextParts = [];
-    const streamId = this.streamId;
-    this.streamId = undefined;
-    this.#cursor = 0;
-    await this.ctx.runMutation(this.component.streams.finish, {
-      streamId,
-      finalDelta,
-    });
   }
+  return compressed;
 }

package/src/client/threads.ts

@@ -0,0 +1,83 @@
+import type { WithoutSystemFields } from "convex/server";
+import type { ThreadDoc } from "../validators.js";
+import type {
+  ActionCtx,
+  AgentComponent,
+  MutationCtx,
+  QueryCtx,
+} from "./types.js";
+
+/**
+ * Create a thread to store messages with an Agent.
+ * @param ctx The context from a mutation or action.
+ * @param component The Agent component, usually `components.agent`.
+ * @param args The associated thread metadata.
+ * @returns The id of the created thread.
+ */
+export async function createThread(
+  ctx: MutationCtx | ActionCtx,
+  component: AgentComponent,
+  args?: { userId?: string | null; title?: string; summary?: string },
+) {
+  const { _id: threadId } = await ctx.runMutation(
+    component.threads.createThread,
+    {
+      userId: args?.userId ?? undefined,
+      title: args?.title,
+      summary: args?.summary,
+    },
+  );
+  return threadId;
+}
+
+/**
+ * Get the metadata for a thread.
+ * @param ctx A ctx object from a query, mutation, or action.
+ * @param args.threadId The thread to get the metadata for.
+ * @returns The metadata for the thread.
+ */
+export async function getThreadMetadata(
+  ctx: QueryCtx | MutationCtx | ActionCtx,
+  component: AgentComponent,
+  args: { threadId: string },
+): Promise<ThreadDoc> {
+  const thread = await ctx.runQuery(component.threads.getThread, {
+    threadId: args.threadId,
+  });
+  if (!thread) {
+    throw new Error("Thread not found");
+  }
+  return thread;
+}
+
+export async function updateThreadMetadata(
+  ctx: MutationCtx | ActionCtx,
+  component: AgentComponent,
+  args: { threadId: string; patch: Partial<WithoutSystemFields<ThreadDoc>> },
+) {
+  return ctx.runMutation(component.threads.updateThread, {
+    threadId: args.threadId,
+    patch: args.patch,
+  });
+}
+
+/**
+ * Search for threads by title, paginated.
+ * @param ctx The context passed from the query/mutation/action.
+ * @returns The threads matching the search, paginated.
+ */
+export async function searchThreadTitles(
+  ctx: QueryCtx | MutationCtx | ActionCtx,
+  component: AgentComponent,
+  {
+    userId,
+    query,
+    limit,
+  }: { userId?: string | undefined; query: string; limit?: number },
+): Promise<ThreadDoc[]> {
+  return ctx.runQuery(component.threads.searchThreadTitles, {
+    userId,
+    query,
+    limit: limit ?? 10,
+  });
+}