@convex-dev/agent 0.2.10 → 0.2.11-alpha.1

package/src/client/index.ts

@@ -89,9 +89,7 @@ import type {
   ObjectSchema,
   Options,
   RawRequestResponseHandler,
-  RunActionCtx,
-  RunMutationCtx,
-  RunQueryCtx,
+  MutationCtx,
   StorageOptions,
   StreamingTextArgs,
   StreamObjectArgs,
@@ -100,10 +98,12 @@ import type {
   Thread,
   UsageHandler,
   UserActionCtx,
+  QueryCtx,
 } from "./types.js";
 
 export { stepCountIs } from "ai";
 export {
+  docsToModelMessages,
   toModelMessage,
   //** @deprecated use toModelMessage instead */
   toModelMessage as deserializeMessage,
@@ -268,7 +268,7 @@ export class Agent<
    * @returns The threadId of the new thread and the thread object.
    */
   async createThread(
-    ctx: RunActionCtx & CustomCtx,
+    ctx: ActionCtx & CustomCtx,
     args?: {
       /**
        * The userId to associate with the thread. If not provided, the thread will be
@@ -295,7 +295,7 @@ export class Agent<
    * @returns The threadId of the new thread.
    */
   async createThread(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx,
     args?: {
       /**
        * The userId to associate with the thread. If not provided, the thread will be
@@ -313,7 +313,7 @@ export class Agent<
     },
   ): Promise<{ threadId: string }>;
   async createThread(
-    ctx: (ActionCtx & CustomCtx) | RunMutationCtx,
+    ctx: (ActionCtx & CustomCtx) | MutationCtx,
     args?: { userId: string | null; title?: string; summary?: string },
   ): Promise<{ threadId: string; thread?: Thread<AgentTools> }> {
     const threadId = await createThread(ctx, this.component, args);
@@ -834,7 +834,7 @@ export class Agent<
    * @returns The messageId of the saved message.
    */
   async saveMessage(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: SaveMessageArgs & {
       /**
        * If true, it will not generate embeddings for the message.
@@ -872,7 +872,7 @@ export class Agent<
    * @returns
    */
   async saveMessages(
-    ctx: RunMutationCtx | RunActionCtx,
+    ctx: MutationCtx | ActionCtx,
     args: SaveMessagesArgs & {
       /**
        * Skip generating embeddings for the messages. Useful if you're
@@ -928,7 +928,7 @@ export class Agent<
    * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
    */
   async listMessages(
-    ctx: RunQueryCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: {
       threadId: string;
       paginationOpts: PaginationOptions;
@@ -948,7 +948,7 @@ export class Agent<
    * @returns The deltas for each stream from their existing cursor.
    */
   async syncStreams(
-    ctx: RunQueryCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: {
       threadId: string;
       streamArgs: StreamArgs | undefined;
@@ -968,7 +968,7 @@ export class Agent<
    * @returns
    */
   async fetchContextMessages(
-    ctx: RunQueryCtx | RunActionCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: {
       userId: string | undefined;
       threadId: string | undefined;
@@ -1030,7 +1030,7 @@ export class Agent<
    * @returns The metadata for the thread.
    */
   async getThreadMetadata(
-    ctx: RunQueryCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: { threadId: string },
   ): Promise<ThreadDoc> {
     return getThreadMetadata(ctx, this.component, args);
@@ -1044,7 +1044,7 @@ export class Agent<
    * @returns The updated thread metadata.
    */
   async updateThreadMetadata(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: {
       threadId: string;
       patch: Partial<
@@ -1065,7 +1065,7 @@ export class Agent<
    * @returns The embeddings for the messages.
    */
   async generateEmbeddings(
-    ctx: RunActionCtx,
+    ctx: ActionCtx,
     args: { userId: string | undefined; threadId: string | undefined },
     messages: (ModelMessage | Message)[],
   ): Promise<
@@ -1091,7 +1091,7 @@ export class Agent<
    * @param args The messageIds to generate embeddings for.
    */
   async generateAndSaveEmbeddings(
-    ctx: RunActionCtx,
+    ctx: ActionCtx,
     args: { messageIds: string[] },
   ) {
     const messages = (
@@ -1251,7 +1251,7 @@ export class Agent<
    *   the generateText call.
    */
   async finalizeMessage(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: {
       messageId: string;
       result: { status: "failed"; error: string } | { status: "success" };
@@ -1269,7 +1269,7 @@ export class Agent<
    * @param args The message fields to update.
    */
   async updateMessage(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: {
       /** The id of the message to update. */
       messageId: string;
@@ -1319,7 +1319,7 @@ export class Agent<
    * @param args The ids of the messages to delete.
    */
   async deleteMessages(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: { messageIds: string[] },
   ): Promise<void> {
     await ctx.runMutation(this.component.messages.deleteByIds, args);
@@ -1332,7 +1332,7 @@ export class Agent<
    * @param args The id of the message to delete.
    */
   async deleteMessage(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: { messageId: string },
   ): Promise<void> {
     await ctx.runMutation(this.component.messages.deleteByIds, {
@@ -1378,7 +1378,7 @@ export class Agent<
    * @param args The range of messages to delete.
    */
   async deleteMessageRange(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: {
       threadId: string;
       startOrder: number;
@@ -1404,7 +1404,7 @@ export class Agent<
    * @param args The id of the thread to delete and optionally the page size to use for the delete.
    */
   async deleteThreadAsync(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: { threadId: string; pageSize?: number },
   ): Promise<void> {
     await ctx.runMutation(this.component.threads.deleteAllForThreadIdAsync, {
@@ -1421,7 +1421,7 @@ export class Agent<
    * @param args The id of the thread to delete and optionally the page size to use for the delete.
    */
   async deleteThreadSync(
-    ctx: RunActionCtx,
+    ctx: ActionCtx,
     args: { threadId: string; pageSize?: number },
   ): Promise<void> {
     await ctx.runAction(this.component.threads.deleteAllForThreadIdSync, {

package/src/client/messages.ts

@@ -12,7 +12,12 @@ import {
 } from "../validators.js";
 import { serializeMessage } from "../mapping.js";
 import { toUIMessages, type UIMessage } from "../UIMessages.js";
-import type { AgentComponent, RunMutationCtx, RunQueryCtx } from "./types.js";
+import type {
+  AgentComponent,
+  MutationCtx,
+  QueryCtx,
+  ActionCtx,
+} from "./types.js";
 import { parse } from "convex-helpers/validators";
 
 /**
@@ -27,7 +32,7 @@ import { parse } from "convex-helpers/validators";
  * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
  */
 export async function listMessages(
-  ctx: RunQueryCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
   {
     threadId,
@@ -58,7 +63,7 @@ export async function listMessages(
 }
 
 export async function listUIMessages(
-  ctx: RunQueryCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
   args: {
     threadId: string;
@@ -105,7 +110,7 @@ export type SaveMessagesArgs = {
  * Explicitly save messages associated with the thread (& user if provided)
  */
 export async function saveMessages(
-  ctx: RunMutationCtx,
+  ctx: MutationCtx,
   component: AgentComponent,
   args: SaveMessagesArgs & {
     /**
@@ -193,7 +198,7 @@ export type SaveMessageArgs = {
  * @returns The messageId of the saved message.
  */
 export async function saveMessage(
-  ctx: RunMutationCtx,
+  ctx: MutationCtx,
   component: AgentComponent,
   args: SaveMessageArgs & {
     /**

package/src/client/saveInputMessages.test.ts

@@ -1,8 +1,12 @@
 import { beforeEach, describe, expect, test, vi } from "vitest";
 import { saveInputMessages } from "./saveInputMessages.js";
 import type { MessageDoc } from "../validators.js";
-import type { RunActionCtx } from "./types.js";
-import { defineSchema } from "convex/server";
+import type { ActionCtx } from "./types.js";
+import {
+  defineSchema,
+  type Auth,
+  type StorageActionWriter,
+} from "convex/server";
 import { initConvexTest } from "./setup.test.js";
 import { components } from "./setup.test.js";
 
@@ -59,7 +63,7 @@ describe("saveInputMessages", () => {
   const mockComponent = components.agent;
 
   let t = initConvexTest(schema);
-  let ctx: RunActionCtx;
+  let ctx: ActionCtx;
 
   beforeEach(() => {
     vi.clearAllMocks();
@@ -68,7 +72,9 @@ describe("saveInputMessages", () => {
       runQuery: t.query,
       runAction: t.action,
       runMutation: t.mutation,
-    } as RunActionCtx;
+      auth: {} as Auth,
+      storage: {} as StorageActionWriter,
+    } as ActionCtx;
 
     mockSaveMessages.mockResolvedValue({
       messages: [
@@ -358,7 +364,9 @@ describe("saveInputMessages", () => {
         runQuery: vi.fn(),
         runMutation: vi.fn(),
         runAction: vi.fn(),
-      } as RunActionCtx;
+        auth: {} as Auth,
+        storage: {} as StorageActionWriter,
+      } as ActionCtx;
 
       await saveInputMessages(actionCtx, mockComponent, {
         ...defaultArgs,

package/src/client/saveInputMessages.ts

@@ -2,10 +2,10 @@ import { type ModelMessage } from "ai";
 import type { MessageDoc } from "../validators.js";
 import { embedMessages, getPromptArray } from "./search.js";
 import type {
+  ActionCtx,
   AgentComponent,
   Config,
-  RunActionCtx,
-  RunMutationCtx,
+  MutationCtx,
 } from "./types.js";
 import { saveMessages } from "./messages.js";
 import type { Message } from "../validators.js";
@@ -13,7 +13,7 @@ import { assert } from "convex-helpers";
 import type { VectorDimension } from "../component/vector/tables.js";
 
 export async function saveInputMessages(
-  ctx: RunMutationCtx | RunActionCtx,
+  ctx: MutationCtx | ActionCtx,
   component: AgentComponent,
   {
     threadId,

package/src/client/search.test.ts

@@ -7,9 +7,14 @@ import {
   type MockedFunction,
 } from "vitest";
 import type { ModelMessage } from "ai";
-import { defineSchema } from "convex/server";
+import {
+  defineSchema,
+  type Auth,
+  type StorageActionWriter,
+  type StorageReader,
+} from "convex/server";
 import type { MessageDoc } from "../validators.js";
-import type { RunActionCtx, RunQueryCtx } from "./types.js";
+import type { ActionCtx, QueryCtx } from "./types.js";
 import {
   fetchContextWithPrompt,
   fetchContextMessages,
@@ -43,8 +48,8 @@ const schema = defineSchema({});
 
 describe("search.ts", () => {
   let t = initConvexTest(schema);
-  let mockCtx: RunActionCtx;
-  let ctx: RunActionCtx;
+  let mockCtx: ActionCtx;
+  let ctx: ActionCtx;
 
   // Shared helper functions
   async function createTestThread(userId: string) {
@@ -88,13 +93,15 @@ describe("search.ts", () => {
       runQuery: t.query,
       runAction: t.action,
       runMutation: t.mutation,
-    } as RunActionCtx;
+    } as ActionCtx;
 
     mockCtx = {
       runQuery: vi.fn(),
       runAction: vi.fn(),
       runMutation: vi.fn(),
-    } satisfies RunActionCtx;
+      auth: {} as Auth,
+      storage: {} as StorageActionWriter,
+    } satisfies ActionCtx;
 
     // Mock process.env to avoid file inlining in tests
     process.env.CONVEX_CLOUD_URL = "https://example.convex.cloud";
@@ -183,10 +190,25 @@ describe("search.ts", () => {
     it("should filter out orphaned tool messages", () => {
       const messages: MessageDoc[] = [
         {
-          _id: "1",
+          _id: "0",
           message: { role: "user", content: "Hello" },
           order: 1,
         } as MessageDoc,
+        {
+          _id: "1",
+          message: {
+            role: "assistant",
+            content: [
+              {
+                type: "tool-call",
+                toolCallId: "call_orphaned",
+                toolName: "test",
+                args: {},
+              },
+            ],
+          },
+          order: 1,
+        } as MessageDoc,
         {
           _id: "2",
           message: {
@@ -194,18 +216,24 @@ describe("search.ts", () => {
             content: [
               {
                 type: "tool-result",
-                toolCallId: "call_orphaned",
+                toolCallId: "result_orphaned",
                 result: "orphaned",
               },
             ],
           },
           order: 2,
         } as MessageDoc,
+        {
+          _id: "3",
+          message: { role: "assistant", content: "I'll help you with that" },
+          order: 1,
+        } as MessageDoc,
       ];
 
       const result = filterOutOrphanedToolMessages(messages);
-      expect(result).toHaveLength(1);
-      expect(result[0]._id).toBe("1");
+      expect(result).toHaveLength(2);
+      expect(result[0]._id).toBe("0");
+      expect(result[1]._id).toBe("3");
     });
   });
 
@@ -227,7 +255,7 @@ describe("search.ts", () => {
       ];
 
       (
-        mockCtx.runQuery as MockedFunction<RunActionCtx["runQuery"]>
+        mockCtx.runQuery as MockedFunction<ActionCtx["runQuery"]>
       ).mockResolvedValue({
         page: mockPage,
       });
@@ -269,7 +297,7 @@ describe("search.ts", () => {
       ];
 
       (
-        mockCtx.runAction as MockedFunction<RunActionCtx["runAction"]>
+        mockCtx.runAction as MockedFunction<ActionCtx["runAction"]>
       ).mockResolvedValue(searchResults);
 
       const result = await fetchContextMessages(mockCtx, components.agent, {
@@ -293,7 +321,8 @@ describe("search.ts", () => {
       const mockQueryCtx = {
         runQuery: vi.fn().mockResolvedValue({ page: [] }),
         // No runAction method
-      } as RunQueryCtx;
+        storage: {} as StorageReader,
+      } as QueryCtx;
 
       await expect(
         fetchContextMessages(mockQueryCtx, components.agent, {
@@ -948,11 +977,7 @@ describe("search.ts", () => {
 
       const contextHandler = vi.fn(async (ctx, args) => {
         // Put existing responses first to test they're properly identified
-        return [
-          ...args.recent,
-          ...args.existingResponses,
-          ...args.inputPrompt,
-        ];
+        return [...args.recent, ...args.existingResponses, ...args.inputPrompt];
       });
 
       const result = await fetchContextWithPrompt(ctx, components.agent, {
@@ -979,7 +1004,7 @@ describe("search.ts", () => {
           inputPrompt: expect.arrayContaining([
             expect.objectContaining({ content: "New replacement prompt" }),
           ]),
-        })
+        }),
       );
 
       expect(result.messages).toHaveLength(4);

package/src/client/search.ts

@@ -20,15 +20,16 @@ import {
 } from "../shared.js";
 import type { Message } from "../validators.js";
 import type {
+  ActionCtx,
   AgentComponent,
   Config,
   ContextOptions,
   Options,
-  RunActionCtx,
-  RunQueryCtx,
+  QueryCtx,
+  MutationCtx,
 } from "./types.js";
 import { inlineMessagesFiles } from "./files.js";
-import { toModelMessage } from "../mapping.js";
+import { docsToModelMessages, toModelMessage } from "../mapping.js";
 
 const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0;
 // 10k characters should be more than enough for most cases, and stays under
@@ -49,7 +50,7 @@ export type GetEmbedding = (text: string) => Promise<{
  * @returns
  */
 export async function fetchContextMessages(
-  ctx: RunQueryCtx | RunActionCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
   args: {
     userId: string | undefined;
@@ -85,7 +86,7 @@ export async function fetchContextMessages(
 }
 
 export async function fetchRecentAndSearchMessages(
-  ctx: RunQueryCtx | RunActionCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
   args: {
     userId: string | undefined;
@@ -229,23 +230,48 @@ export async function fetchRecentAndSearchMessages(
  */
 export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   const toolCallIds = new Set<string>();
+  const toolResultIds = new Set<string>();
   const result: MessageDoc[] = [];
   for (const doc of docs) {
-    if (
-      doc.message?.role === "assistant" &&
-      Array.isArray(doc.message.content)
-    ) {
+    if (doc.message && Array.isArray(doc.message.content)) {
       for (const content of doc.message.content) {
         if (content.type === "tool-call") {
           toolCallIds.add(content.toolCallId);
+        } else if (content.type === "tool-result") {
+          toolResultIds.add(content.toolCallId);
         }
       }
-      result.push(doc);
+    }
+  }
+  for (const doc of docs) {
+    if (
+      doc.message?.role === "assistant" &&
+      Array.isArray(doc.message.content)
+    ) {
+      const content = doc.message.content.filter(
+        (p) => p.type !== "tool-call" || toolResultIds.has(p.toolCallId),
+      );
+      if (content.length) {
+        result.push({
+          ...doc,
+          message: {
+            ...doc.message,
+            content,
+          },
+        });
+      }
     } else if (doc.message?.role === "tool") {
-      if (doc.message.content.every((c) => toolCallIds.has(c.toolCallId))) {
-        result.push(doc);
-      } else {
-        console.debug("Filtering out orphaned tool message", doc);
+      const content = doc.message.content.filter((c) =>
+        toolCallIds.has(c.toolCallId),
+      );
+      if (content.length) {
+        result.push({
+          ...doc,
+          message: {
+            ...doc.message,
+            content,
+          },
+        });
       }
     } else {
       result.push(doc);
@@ -259,7 +285,7 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
  * This will not save the embeddings to the database.
  */
 export async function embedMessages(
-  ctx: RunActionCtx,
+  ctx: ActionCtx,
   {
     userId,
     threadId,
@@ -328,7 +354,7 @@ export async function embedMessages(
  * @returns The embeddings for the strings, matching the order of the values.
  */
 export async function embedMany(
-  ctx: RunActionCtx,
+  ctx: ActionCtx,
   {
     userId,
     threadId,
@@ -386,7 +412,7 @@ export async function embedMany(
  * @param messages The messages to embed, in the Agent MessageDoc format.
  */
 export async function generateAndSaveEmbeddings(
-  ctx: RunActionCtx,
+  ctx: ActionCtx,
   component: AgentComponent,
   args: {
     threadId: string | undefined;
@@ -431,7 +457,7 @@ export async function generateAndSaveEmbeddings(
  * promptMessageId message.
  */
 export async function fetchContextWithPrompt(
-  ctx: RunActionCtx,
+  ctx: ActionCtx,
   component: AgentComponent,
   args: {
     prompt: string | (ModelMessage | Message)[] | undefined;
@@ -522,28 +548,19 @@ export async function fetchContextWithPrompt(
     }
   }
 
-  const search = searchMessages
-    .map((m) => m.message)
-    .filter((m) => !!m)
-    .map(toModelMessage);
-  const recent = prePromptDocs
-    .map((m) => m.message)
-    .filter((m) => !!m)
-    .map(toModelMessage);
+  const search = docsToModelMessages(searchMessages);
+  const recent = docsToModelMessages(prePromptDocs);
   const inputMessages = messages.map(toModelMessage);
   const inputPrompt = promptArray.map(toModelMessage);
-  const existingResponses = existingResponseDocs
-    .map((m) => m.message)
-    .filter((m) => !!m)
-    .map(toModelMessage);
+  const existingResponses = docsToModelMessages(existingResponseDocs);
 
-    const allMessages = [
-      ...search,
-      ...recent,
-      ...inputMessages,
-      ...inputPrompt,
-      ...existingResponses,
-    ];
+  const allMessages = [
+    ...search,
+    ...recent,
+    ...inputMessages,
+    ...inputPrompt,
+    ...existingResponses,
+  ];
   let processedMessages = args.contextHandler
     ? await args.contextHandler(ctx, {
         allMessages,

package/src/client/streaming.ts

@@ -19,9 +19,10 @@ import {
   type StreamMessage,
 } from "../validators.js";
 import type {
+  ActionCtx,
   AgentComponent,
-  RunMutationCtx,
-  RunQueryCtx,
+  MutationCtx,
+  QueryCtx,
   SyncStreamsReturnValue,
 } from "./types.js";
 
@@ -45,7 +46,7 @@ export const vStreamMessagesReturnValue = v.object({
  * @returns The deltas for each stream from their existing cursor.
  */
 export async function syncStreams(
-  ctx: RunQueryCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
   {
     threadId,
@@ -80,7 +81,7 @@ export async function syncStreams(
 }
 
 export async function abortStream(
-  ctx: RunMutationCtx,
+  ctx: MutationCtx,
   component: AgentComponent,
   args: { reason: string } & (
     | { streamId: string }
@@ -111,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,
@@ -212,7 +213,7 @@ export class DeltaStreamer<T> {
 
   constructor(
     public readonly component: AgentComponent,
-    public readonly ctx: RunMutationCtx,
+    public readonly ctx: MutationCtx,
     config: {
       throttleMs: number | undefined;
       onAsyncAbort: (reason: string) => Promise<void>;