@convex-dev/agent 0.2.4 → 0.2.5

package/src/client/mockModel.ts

@@ -0,0 +1,195 @@
+import type {
+  LanguageModelV2,
+  LanguageModelV2StreamPart,
+} from "@ai-sdk/provider";
+import type { ReasoningPart, TextPart } from "@ai-sdk/provider-utils";
+import { simulateReadableStream } from "ai";
+
+const longDefaultText = `
+A A A A A A A A A A A A A A A
+B B B B B B B B B B B B B B B
+C C C C C C C C C C C C C C C
+D D D D D D D D D D D D D D D
+`;
+const defaultUsage = { outputTokens: 10, inputTokens: 3, totalTokens: 13 };
+
+export type MockModelArgs = {
+  provider?: LanguageModelV2["provider"];
+  modelId?: LanguageModelV2["modelId"];
+  supportedUrls?:
+    | LanguageModelV2["supportedUrls"]
+    | (() => LanguageModelV2["supportedUrls"]);
+  chunkDelayInMs?: number;
+  initialDelayInMs?: number;
+  // provide either content or doGenerate & doStream
+  content?: (TextPart | ReasoningPart)[];
+  doGenerate?: LanguageModelV2["doGenerate"];
+  doStream?: LanguageModelV2["doStream"];
+  fail?:
+    | boolean
+    | {
+        probability?: number;
+        error?: string;
+      };
+};
+
+export function mockModel(args?: MockModelArgs): LanguageModelV2 {
+  return new MockLanguageModel(args ?? {});
+}
+
+export class MockLanguageModel implements LanguageModelV2 {
+  readonly specificationVersion = "v2";
+
+  private _supportedUrls: () => LanguageModelV2["supportedUrls"];
+
+  readonly provider: LanguageModelV2["provider"];
+  readonly modelId: LanguageModelV2["modelId"];
+
+  doGenerate: LanguageModelV2["doGenerate"];
+  doStream: LanguageModelV2["doStream"];
+
+  doGenerateCalls: Parameters<LanguageModelV2["doGenerate"]>[0][] = [];
+  doStreamCalls: Parameters<LanguageModelV2["doStream"]>[0][] = [];
+
+  constructor(args: MockModelArgs) {
+    this.provider = args.provider || "mock-provider";
+    this.modelId = args.modelId || "mock-model-id";
+    const {
+      content = [{ type: "text", text: longDefaultText }],
+      chunkDelayInMs = 200,
+      initialDelayInMs = 1000,
+      supportedUrls = {},
+    } = args;
+    const fail =
+      args.fail &&
+      (args.fail === true ||
+        !args.fail.probability ||
+        Math.random() < args.fail.probability);
+    const error =
+      (typeof args.fail === "object" && args.fail.error) ||
+      "Mock error message";
+
+    const chunks: LanguageModelV2StreamPart[] = [
+      { type: "stream-start", warnings: [] },
+    ];
+    chunks.push(
+      ...content.flatMap((c, ci) => {
+        const deltas = c.text.split(" ");
+        let parts: LanguageModelV2StreamPart[] = [];
+        if (c.type === "reasoning") {
+          parts.push({
+            type: "reasoning-start",
+            id: `${ci}-reasoning-start`,
+          });
+          parts.push(
+            ...deltas.map(
+              (delta, di) =>
+                ({
+                  type: "reasoning-delta",
+                  delta,
+                  id: `${ci}-reasoning-${di}`,
+                  providerMetadata: {
+                    mockProvider: { mock: { reasoningDetails: null } },
+                  },
+                }) satisfies LanguageModelV2StreamPart,
+            ),
+          );
+          parts.push({
+            type: "reasoning-end",
+            id: `${ci}-reasoning-end`,
+          });
+        } else if (c.type === "text") {
+          parts.push({
+            type: "text-start",
+            id: `${ci}-text-start`,
+          });
+          parts = deltas.map((delta, di) => ({
+            type: "text-delta",
+            delta,
+            id: `${ci}-text-${di}`,
+          }));
+          parts.push({
+            type: "text-end",
+            id: `${ci}-text-end`,
+          });
+        }
+        return parts;
+      }),
+    );
+    if (fail) {
+      chunks.push({
+        type: "error",
+        error,
+      });
+    }
+    chunks.push({
+      type: "finish",
+      finishReason: fail ? "error" : "stop",
+      usage: defaultUsage,
+      providerMetadata: {
+        mockProvider: { mock: "mock metadata" },
+      },
+    });
+    this.doGenerate = async (options) => {
+      this.doGenerateCalls.push(options);
+
+      if (fail) {
+        throw new Error(error);
+      }
+      if (typeof args.doGenerate === "function") {
+        return args.doGenerate(options);
+      } else if (Array.isArray(args.doGenerate)) {
+        return args.doGenerate[this.doGenerateCalls.length];
+      } else if (content) {
+        return {
+          content,
+          finishReason: "stop",
+          usage: defaultUsage,
+          providerMetadata: { mockProvider: { mock: "mock metadata" } },
+          warnings: [],
+        };
+      } else {
+        throw new Error("Unexpected: no content or doGenerate");
+      }
+    };
+    this._supportedUrls =
+      typeof supportedUrls === "function"
+        ? supportedUrls
+        : async () => supportedUrls;
+    this.doStream = async (options) => {
+      this.doStreamCalls.push(options);
+
+      if (typeof args.doStream === "function") {
+        return args.doStream(options);
+      } else if (Array.isArray(args.doStream)) {
+        return args.doStream[this.doStreamCalls.length];
+      } else if (content) {
+        const stream = simulateReadableStream({
+          chunks,
+          initialDelayInMs,
+          chunkDelayInMs,
+        });
+        if (options.abortSignal) {
+          throw new Error("abortSignal in mock model");
+        }
+        return {
+          stream,
+          request: { body: {} },
+          response: { headers: {} },
+        };
+      } else if (args.doStream) {
+        return args.doStream;
+      } else {
+        throw new Error("Provide either content or doStream");
+      }
+    };
+    this._supportedUrls =
+      typeof supportedUrls === "function"
+        ? supportedUrls
+        : async () => supportedUrls;
+  }
+
+  get supportedUrls() {
+    return this._supportedUrls();
+  }
+}

package/src/client/streaming.ts

@@ -197,6 +197,7 @@ export class DeltaStreamer {
     public readonly component: AgentComponent,
     public readonly ctx: RunActionCtx,
     options: true | StreamingOptions,
+    private onAsyncAbort: (reason: string) => Promise<void>,
     public readonly metadata: {
       threadId: string;
       userId?: string;
@@ -210,21 +211,22 @@ export class DeltaStreamer {
     },
   ) {
     this.options =
-      typeof options === "boolean"
+      options === true
         ? DEFAULT_STREAMING_OPTIONS
         : { ...DEFAULT_STREAMING_OPTIONS, ...options };
     this.#nextParts = [];
     this.abortController = new AbortController();
     if (metadata.abortSignal) {
       metadata.abortSignal.addEventListener("abort", async () => {
+        if (this.abortController.signal.aborted) {
+          return;
+        }
         if (this.streamId) {
           this.abortController.abort();
-          const finalDelta = this.#createDelta();
           await this.#ongoingWrite;
           await this.ctx.runMutation(this.component.streams.abort, {
             streamId: this.streamId,
             reason: "abortSignal",
-            finalDelta,
           });
         }
       });
@@ -265,9 +267,12 @@ export class DeltaStreamer {
         delta,
       );
       if (!success) {
+        await this.onAsyncAbort("async abort");
         this.abortController.abort();
+        return;
       }
     } catch (e) {
+      await this.onAsyncAbort(e instanceof Error ? e.message : "unknown error");
       this.abortController.abort();
       throw e;
     }
@@ -302,11 +307,9 @@ export class DeltaStreamer {
     if (!this.streamId) {
       return;
     }
-    const finalDelta = this.#createDelta();
     await this.#ongoingWrite;
     await this.ctx.runMutation(this.component.streams.finish, {
       streamId: this.streamId,
-      finalDelta,
     });
   }
 
@@ -318,12 +321,10 @@ export class DeltaStreamer {
     if (!this.streamId) {
       return;
     }
-    const finalDelta = this.#createDelta();
     await this.#ongoingWrite;
     await this.ctx.runMutation(this.component.streams.abort, {
       streamId: this.streamId,
       reason,
-      finalDelta,
     });
   }
 }

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

@@ -260,6 +260,31 @@ export type Mounts = {
                           toolName: string;
                           type: "tool-result";
                         }
+                      | {
+                          id: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "url";
+                          title: string;
+                          type: "source";
+                          url: string;
+                        }
+                      | {
+                          filename?: string;
+                          id: string;
+                          mediaType: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "document";
+                          title: string;
+                          type: "source";
+                        }
                     >;
                 providerOptions?: Record<string, Record<string, any>>;
                 role: "assistant";
@@ -469,6 +494,31 @@ export type Mounts = {
                           toolName: string;
                           type: "tool-result";
                         }
+                      | {
+                          id: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "url";
+                          title: string;
+                          type: "source";
+                          url: string;
+                        }
+                      | {
+                          filename?: string;
+                          id: string;
+                          mediaType: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "document";
+                          title: string;
+                          type: "source";
+                        }
                     >;
                 providerOptions?: Record<string, Record<string, any>>;
                 role: "assistant";
@@ -687,6 +737,25 @@ export type Mounts = {
                         toolName: string;
                         type: "tool-result";
                       }
+                    | {
+                        id: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "url";
+                        title: string;
+                        type: "source";
+                        url: string;
+                      }
+                    | {
+                        filename?: string;
+                        id: string;
+                        mediaType: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "document";
+                        title: string;
+                        type: "source";
+                      }
                   >;
               providerOptions?: Record<string, Record<string, any>>;
               role: "assistant";
@@ -918,6 +987,31 @@ export type Mounts = {
                           toolName: string;
                           type: "tool-result";
                         }
+                      | {
+                          id: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "url";
+                          title: string;
+                          type: "source";
+                          url: string;
+                        }
+                      | {
+                          filename?: string;
+                          id: string;
+                          mediaType: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "document";
+                          title: string;
+                          type: "source";
+                        }
                     >;
                 providerOptions?: Record<string, Record<string, any>>;
                 role: "assistant";
@@ -1121,6 +1215,25 @@ export type Mounts = {
                         toolName: string;
                         type: "tool-result";
                       }
+                    | {
+                        id: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "url";
+                        title: string;
+                        type: "source";
+                        url: string;
+                      }
+                    | {
+                        filename?: string;
+                        id: string;
+                        mediaType: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "document";
+                        title: string;
+                        type: "source";
+                      }
                   >;
               providerOptions?: Record<string, Record<string, any>>;
               role: "assistant";
@@ -1317,6 +1430,25 @@ export type Mounts = {
                         toolName: string;
                         type: "tool-result";
                       }
+                    | {
+                        id: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "url";
+                        title: string;
+                        type: "source";
+                        url: string;
+                      }
+                    | {
+                        filename?: string;
+                        id: string;
+                        mediaType: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "document";
+                        title: string;
+                        type: "source";
+                      }
                   >;
               providerOptions?: Record<string, Record<string, any>>;
               role: "assistant";
@@ -1527,6 +1659,31 @@ export type Mounts = {
                           toolName: string;
                           type: "tool-result";
                         }
+                      | {
+                          id: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "url";
+                          title: string;
+                          type: "source";
+                          url: string;
+                        }
+                      | {
+                          filename?: string;
+                          id: string;
+                          mediaType: string;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          sourceType: "document";
+                          title: string;
+                          type: "source";
+                        }
                     >;
                 providerOptions?: Record<string, Record<string, any>>;
                 role: "assistant";
@@ -1661,6 +1818,25 @@ export type Mounts = {
                         toolName: string;
                         type: "tool-result";
                       }
+                    | {
+                        id: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "url";
+                        title: string;
+                        type: "source";
+                        url: string;
+                      }
+                    | {
+                        filename?: string;
+                        id: string;
+                        mediaType: string;
+                        providerMetadata?: Record<string, Record<string, any>>;
+                        providerOptions?: Record<string, Record<string, any>>;
+                        sourceType: "document";
+                        title: string;
+                        type: "source";
+                      }
                   >;
               providerOptions?: Record<string, Record<string, any>>;
               role: "assistant";

package/src/component/messages.ts

@@ -351,7 +351,7 @@ export const finalizeMessage = mutation({
       return;
     }
     // See if we can add any in-progress data
-    if (message.message === undefined) {
+    if (!message.message?.content.length) {
       const messages = await getStreamingMessagesWithMetadata(
         ctx,
         message,

package/src/mapping.ts

@@ -20,6 +20,7 @@ import {
 } from "ai";
 import {
   vMessageWithMetadata,
+  type vSourcePart,
   type Message,
   type MessageWithMetadata,
   type Usage,
@@ -31,6 +32,7 @@ import {
   type vTextPart,
   type vToolCallPart,
   type vToolResultPart,
+  type SourcePart,
 } from "./validators.js";
 import type { ActionCtx, AgentComponent } from "./client/types.js";
 import type { RunMutationCtx } from "./client/types.js";
@@ -335,6 +337,9 @@ export async function serializeContent(
             ...metadata,
           } satisfies Infer<typeof vRedactedReasoningPart>;
         }
+        case "source": {
+          return part satisfies Infer<typeof vSourcePart>;
+        }
         default:
           return part satisfies Infer<typeof vContent>;
       }
@@ -431,6 +436,8 @@ export function deserializeContent(
               }
             : undefined,
         } satisfies ReasoningPart;
+      case "source":
+        return part satisfies SourcePart;
       default:
         return part satisfies Content;
     }

package/src/react/toUIMessages.ts

@@ -12,8 +12,9 @@ import type {
 } from "ai";
 import { extractText, type MessageDoc } from "../client/index.js";
 import { deserializeMessage, toUIFilePart } from "../mapping.js";
-import type { MessageStatus } from "../validators.js";
+import type { MessageStatus, SourcePart, vSource } from "../validators.js";
 import { sorted } from "../shared.js";
+import type { Infer } from "convex/values";
 
 export type UIMessage<
   METADATA = unknown,
@@ -388,29 +389,23 @@ function createAssistantUIMessage<
           }
           break;
         }
+        default: {
+          const maybeSource = contentPart as SourcePart;
+          if (maybeSource.type === "source") {
+            allParts.push(toSourcePart(maybeSource));
+          } else {
+            console.warn(
+              "Unknown content part type for assistant",
+              contentPart,
+            );
+          }
+        }
       }
     }
 
     // Add source parts
     for (const source of message.sources ?? []) {
-      if (source.sourceType === "url") {
-        allParts.push({
-          type: "source-url",
-          url: source.url!,
-          sourceId: source.id,
-          providerMetadata: message.providerMetadata,
-          title: source.title,
-        } satisfies SourceUrlUIPart);
-      } else {
-        allParts.push({
-          type: "source-document",
-          mediaType: source.mediaType,
-          sourceId: source.id,
-          title: source.title,
-          filename: source.filename,
-          providerMetadata: message.providerMetadata,
-        } satisfies SourceDocumentUIPart);
-      }
+      allParts.push(toSourcePart(source));
     }
   }
 
@@ -422,3 +417,25 @@ function createAssistantUIMessage<
     parts: allParts,
   };
 }
+
+function toSourcePart(
+  part: SourcePart | Infer<typeof vSource>,
+): SourceUrlUIPart | SourceDocumentUIPart {
+  if (part.sourceType === "url") {
+    return {
+      type: "source-url",
+      url: part.url,
+      sourceId: part.id,
+      providerMetadata: part.providerMetadata,
+      title: part.title,
+    } satisfies SourceUrlUIPart;
+  }
+  return {
+    type: "source-document",
+    mediaType: part.mediaType,
+    sourceId: part.id,
+    title: part.title,
+    filename: part.filename,
+    providerMetadata: part.providerMetadata,
+  } satisfies SourceDocumentUIPart;
+}

package/src/validators.ts

@@ -87,6 +87,29 @@ export const vReasoningDetails = v.array(
   ),
 );
 
+export const vSourcePart = v.union(
+  v.object({
+    type: v.literal("source"),
+    sourceType: v.literal("url"),
+    id: v.string(),
+    url: v.string(),
+    title: v.string(),
+    providerOptions,
+    providerMetadata,
+  }),
+  v.object({
+    type: v.literal("source"),
+    sourceType: v.literal("document"),
+    id: v.string(),
+    mediaType: v.string(),
+    title: v.string(),
+    filename: v.optional(v.string()),
+    providerOptions,
+    providerMetadata,
+  }),
+);
+export type SourcePart = Infer<typeof vSourcePart>;
+
 export const vToolCallPart = v.object({
   type: v.literal("tool-call"),
   toolCallId: v.string(),
@@ -136,6 +159,7 @@ export const vAssistantContent = v.union(
       vRedactedReasoningPart,
       vToolCallPart,
       vToolResultPart,
+      vSourcePart,
     ),
   ),
 );