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

package/dist/component/messages.js

@@ -1,15 +1,16 @@
-import { assert, omit } from "convex-helpers";
+import { assert, omit, pick } from "convex-helpers";
 import { mergedStream, stream } from "convex-helpers/server/stream";
-import { paginationOptsValidator } from "convex/server";
-import { DEFAULT_MESSAGE_RANGE, DEFAULT_RECENT_MESSAGES, extractText, isTool, } from "../shared.js";
-import { vMessageEmbeddings, vMessageStatus, vMessageWithMetadataInternal, vPaginationResult, } from "../validators.js";
+import { paginationOptsValidator, } from "convex/server";
+import { DEFAULT_MESSAGE_RANGE, DEFAULT_RECENT_MESSAGES, extractText, isTool, sorted, } from "../shared.js";
+import { vMessageDoc, vMessageEmbeddingsWithDimension, vMessageStatus, vMessageWithMetadataInternal, vPaginationResult, } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
-import { action, internalQuery, mutation, query, } from "./_generated/server.js";
-import { schema, v, vMessageDoc } from "./schema.js";
-import { getThread as _getThread, listThreadsByUserId as _listThreadsByUserId, updateThread as _updateThread, } from "./threads.js";
+import { action, internalMutation, internalQuery, mutation, query, } from "./_generated/server.js";
+import { schema, v } from "./schema.js";
 import { insertVector, searchVectors } from "./vector/index.js";
-import { VectorDimensions, vVectorId, } from "./vector/tables.js";
+import { validateVectorDimension, vVectorId, } from "./vector/tables.js";
 import { changeRefcount } from "./files.js";
+import { getStreamingMessagesWithMetadata } from "./streams.js";
+import { partial } from "convex-helpers/validators";
 function publicMessage(message) {
     return omit(message, ["parentMessageId", "stepId", "files"]);
 }
@@ -23,9 +24,7 @@ export async function deleteMessage(ctx, messageDoc) {
     }
 }
 export const deleteByIds = mutation({
-    args: {
-        messageIds: v.array(v.id("messages")),
-    },
+    args: { messageIds: v.array(v.id("messages")) },
     returns: v.array(v.id("messages")),
     handler: async (ctx, args) => {
         const deletedMessageIds = await Promise.all(args.messageIds.map(async (id) => {
@@ -54,7 +53,12 @@ export const deleteByOrder = mutation({
         lastStepOrder: v.optional(v.number()),
     }),
     handler: async (ctx, args) => {
-        const messages = await orderedMessagesStream(ctx, args.threadId, "asc", args.startOrder)
+        const messages = await orderedMessagesStream(ctx, {
+            threadId: args.threadId,
+            sortOrder: "asc",
+            startOrder: args.startOrder,
+            startOrderBound: "gte",
+        })
             .narrow({
             lowerBound: args.startStepOrder
                 ? [args.startOrder, args.startStepOrder]
@@ -80,16 +84,18 @@ const addMessagesArgs = {
     promptMessageId: v.optional(v.id("messages")),
     agentName: v.optional(v.string()),
     messages: v.array(vMessageWithMetadataInternal),
-    embeddings: v.optional(vMessageEmbeddings),
-    pending: v.optional(v.boolean()),
+    embeddings: v.optional(vMessageEmbeddingsWithDimension),
     failPendingSteps: v.optional(v.boolean()),
+    // A pending message to update. If the pending message failed, abort.
+    pendingMessageId: v.optional(v.id("messages")),
+    // if set to true, these messages will not show up in text or vector search
+    // results for the userId
+    hideFromUserIdSearch: v.optional(v.boolean()),
 };
 export const addMessages = mutation({
     args: addMessagesArgs,
     handler: addMessagesHandler,
-    returns: v.object({
-        messages: v.array(vMessageDoc),
-    }),
+    returns: v.object({ messages: v.array(vMessageDoc) }),
 });
 async function addMessagesHandler(ctx, args) {
     let userId = args.userId;
@@ -99,34 +105,47 @@ async function addMessagesHandler(ctx, args) {
         assert(thread, `Thread ${args.threadId} not found`);
         userId = thread.userId;
     }
-    const { embeddings, failPendingSteps, pending, messages, promptMessageId, ...rest } = args;
-    const parentMessage = promptMessageId && (await ctx.db.get(promptMessageId));
+    const { embeddings, failPendingSteps, messages, promptMessageId, pendingMessageId, hideFromUserIdSearch, ...rest } = args;
+    const promptMessage = promptMessageId && (await ctx.db.get(promptMessageId));
     if (failPendingSteps) {
         assert(args.threadId, "threadId is required to fail pending steps");
         const pendingMessages = await ctx.db
             .query("messages")
             .withIndex("threadId_status_tool_order_stepOrder", (q) => q.eq("threadId", threadId).eq("status", "pending"))
-            .collect();
+            .order("desc")
+            .take(100);
         await Promise.all(pendingMessages
-            .filter((m) => !parentMessage || m.order === parentMessage.order)
-            .map((m) => ctx.db.patch(m._id, { status: "failed", error: "Restarting" })));
+            .filter((m) => !promptMessage || m.order === promptMessage.order)
+            .filter((m) => !pendingMessageId || m._id !== pendingMessageId)
+            .map(async (m) => {
+            if (m.embeddingId) {
+                await ctx.db.delete(m.embeddingId);
+            }
+            await ctx.db.patch(m._id, {
+                status: "failed",
+                error: "Restarting",
+                embeddingId: undefined,
+            });
+        }));
     }
     let order, stepOrder;
     let fail = false;
+    let error;
     if (promptMessageId) {
-        assert(parentMessage, `Parent message ${promptMessageId} not found`);
-        if (parentMessage.status === "failed") {
+        assert(promptMessage, `Parent message ${promptMessageId} not found`);
+        if (promptMessage.status === "failed") {
             fail = true;
+            error = promptMessage.error ?? error ?? "The prompt message failed";
         }
-        order = parentMessage.order;
+        order = promptMessage.order;
         // Defend against there being existing messages with this parent.
         const maxMessage = await getMaxMessage(ctx, threadId, order);
-        stepOrder = maxMessage?.stepOrder ?? parentMessage.stepOrder;
+        stepOrder = maxMessage?.stepOrder ?? promptMessage.stepOrder;
     }
     else {
         const maxMessage = await getMaxMessage(ctx, threadId);
-        order = maxMessage ? maxMessage.order + 1 : 0;
-        stepOrder = -1;
+        order = maxMessage?.order ?? -1;
+        stepOrder = maxMessage?.stepOrder ?? -1;
     }
     const toReturn = [];
     if (embeddings) {
@@ -135,35 +154,75 @@ async function addMessagesHandler(ctx, args) {
     for (let i = 0; i < messages.length; i++) {
         const message = messages[i];
         let embeddingId;
-        if (embeddings && embeddings.vectors[i]) {
+        if (embeddings &&
+            embeddings.vectors[i] &&
+            !fail &&
+            message.status !== "failed") {
             embeddingId = await insertVector(ctx, embeddings.dimension, {
                 vector: embeddings.vectors[i],
                 model: embeddings.model,
                 table: "messages",
-                userId,
+                userId: hideFromUserIdSearch ? undefined : userId,
                 threadId,
             });
         }
-        stepOrder++;
-        const messageId = await ctx.db.insert("messages", {
+        const messageDoc = {
             ...rest,
             ...message,
             embeddingId,
             parentMessageId: promptMessageId,
             userId,
-            order,
             tool: isTool(message.message),
-            text: extractText(message.message),
-            status: fail ? "failed" : pending ? "pending" : "success",
-            error: fail ? "Parent message failed" : undefined,
+            text: hideFromUserIdSearch ? undefined : extractText(message.message),
+            status: fail ? "failed" : (message.status ?? "success"),
+            error: fail ? error : message.error,
+        };
+        // If there is a pending message, we replace that one with the first message
+        // and subsequent ones will follow the regular order/subOrder advancement.
+        if (i === 0 && pendingMessageId) {
+            const pendingMessage = await ctx.db.get(pendingMessageId);
+            assert(pendingMessage, `Pending msg ${pendingMessageId} not found`);
+            if (pendingMessage.status === "failed") {
+                fail = true;
+                error =
+                    `Trying to update a message that failed: ${pendingMessageId}, ` +
+                        `error: ${pendingMessage.error ?? error}`;
+                messageDoc.status = "failed";
+                messageDoc.error = error;
+            }
+            if (message.fileIds) {
+                await changeRefcount(ctx, pendingMessage.fileIds ?? [], message.fileIds);
+            }
+            await ctx.db.replace(pendingMessage._id, {
+                ...messageDoc,
+                order: pendingMessage.order,
+                stepOrder: pendingMessage.stepOrder,
+            });
+            toReturn.push(pendingMessage);
+            continue;
+        }
+        if (message.message.role === "user") {
+            if (promptMessage && promptMessage.order === order) {
+                // see if there's a later message than the parent message order
+                const maxMessage = await getMaxMessage(ctx, threadId);
+                order = (maxMessage?.order ?? order) + 1;
+            }
+            else {
+                order++;
+            }
+            stepOrder = 0;
+        }
+        else {
+            if (order < 0) {
+                order = 0;
+            }
+            stepOrder++;
+        }
+        const messageId = await ctx.db.insert("messages", {
+            ...messageDoc,
+            order,
             stepOrder,
         });
-        // Let's just not set the id field and have it set only in explicit cases.
-        // if (!message.id) {
-        //   await ctx.db.patch(messageId, {
-        //     id: messageId,
-        //   });
-        // }
         if (message.fileIds) {
             await changeRefcount(ctx, [], message.fileIds);
         }
@@ -174,60 +233,90 @@ async function addMessagesHandler(ctx, args) {
 }
 // exported for tests
 export async function getMaxMessage(ctx, threadId, order) {
-    return orderedMessagesStream(ctx, threadId, "desc", order).first();
+    return orderedMessagesStream(ctx, {
+        threadId,
+        sortOrder: "desc",
+        startOrder: order,
+        startOrderBound: "eq",
+    }).first();
 }
-function orderedMessagesStream(ctx, threadId, sortOrder, order) {
+function orderedMessagesStream(ctx, args) {
     return mergedStream([true, false].flatMap((tool) => messageStatuses.map((status) => stream(ctx.db, schema)
         .query("messages")
         .withIndex("threadId_status_tool_order_stepOrder", (q) => {
         const qq = q
-            .eq("threadId", threadId)
+            .eq("threadId", args.threadId)
             .eq("status", status)
             .eq("tool", tool);
-        if (order) {
-            return qq.eq("order", order);
+        if (args.startOrder !== undefined) {
+            if (args.startOrderBound === "gte") {
+                return qq.gte("order", args.startOrder);
+            }
+            else {
+                return qq.eq("order", args.startOrder);
+            }
         }
         return qq;
     })
-        .order(sortOrder))), ["order", "stepOrder"]);
+        .order(args.sortOrder))), ["order", "stepOrder"]);
 }
-export const rollbackMessage = mutation({
+export const finalizeMessage = mutation({
     args: {
         messageId: v.id("messages"),
-        error: v.optional(v.string()),
+        result: v.union(v.object({ status: v.literal("success") }), v.object({ status: v.literal("failed"), error: v.string() })),
     },
     returns: v.null(),
-    handler: async (ctx, { messageId, error }) => {
+    handler: async (ctx, { messageId, result }) => {
         const message = await ctx.db.get(messageId);
         assert(message, `Message ${messageId} not found`);
-        const messages = await orderedMessagesStream(ctx, message.threadId, "asc", message.order).collect();
-        for (const m of messages) {
-            if (m.status === "pending") {
-                await ctx.db.patch(m._id, { status: "failed", error });
+        if (message.status !== "pending") {
+            console.debug("Trying to finalize a message that's already", message.status);
+            return;
+        }
+        // See if we can add any in-progress data
+        if (!message.message?.content.length) {
+            const messages = await getStreamingMessagesWithMetadata(ctx, message, result);
+            if (messages.length > 0) {
+                await addMessagesHandler(ctx, {
+                    messages,
+                    threadId: message.threadId,
+                    agentName: message.agentName,
+                    failPendingSteps: false,
+                    pendingMessageId: messageId,
+                    userId: message.userId,
+                    embeddings: undefined,
+                });
+                return;
             }
         }
-        await ctx.db.patch(messageId, {
-            status: "failed",
-            error: error,
-        });
-    },
-});
-export const commitMessage = mutation({
-    args: {
-        messageId: v.id("messages"),
+        if (result.status === "failed") {
+            if (message.embeddingId) {
+                await ctx.db.delete(message.embeddingId);
+            }
+            await ctx.db.patch(messageId, {
+                status: "failed",
+                error: result.error,
+                embeddingId: undefined,
+            });
+        }
+        else {
+            await ctx.db.patch(messageId, { status: "success" });
+        }
     },
-    returns: v.null(),
-    handler: commitMessageHandler,
 });
 export const updateMessage = mutation({
     args: {
         messageId: v.id("messages"),
-        patch: v.object({
-            message: v.optional(vMessageDoc.fields.message),
-            fileIds: v.optional(v.array(v.id("files"))),
-            status: v.optional(vMessageStatus),
-            error: v.optional(v.string()),
-        }),
+        patch: v.object(partial(pick(schema.tables.messages.validator.fields, [
+            "message",
+            "fileIds",
+            "status",
+            "error",
+            "model",
+            "provider",
+            "providerOptions",
+            "finishReason",
+        ]))),
     },
     returns: vMessageDoc,
     handler: async (ctx, args) => {
@@ -236,83 +325,185 @@ export const updateMessage = mutation({
         if (args.patch.fileIds) {
             await changeRefcount(ctx, message.fileIds ?? [], args.patch.fileIds);
         }
-        const patch = {
-            ...args.patch,
-        };
+        const patch = { ...args.patch };
         if (args.patch.message !== undefined) {
             patch.message = args.patch.message;
             patch.tool = isTool(args.patch.message);
             patch.text = extractText(args.patch.message);
         }
+        if (args.patch.status === "failed") {
+            if (message.embeddingId) {
+                await ctx.db.delete(message.embeddingId);
+            }
+            patch.embeddingId = undefined;
+        }
         await ctx.db.patch(args.messageId, patch);
         return publicMessage((await ctx.db.get(args.messageId)));
     },
 });
-async function commitMessageHandler(ctx, { messageId }) {
-    const message = await ctx.db.get(messageId);
-    assert(message, `Message ${messageId} not found`);
-    const order = message.order;
-    const messages = await mergedStream([true, false].map((tool) => stream(ctx.db, schema)
-        .query("messages")
-        .withIndex("threadId_status_tool_order_stepOrder", (q) => q
-        .eq("threadId", message.threadId)
-        .eq("status", "pending")
-        .eq("tool", tool)
-        .eq("order", order))), ["order", "stepOrder"]).collect();
-    for (const message of messages) {
-        await ctx.db.patch(message._id, { status: "success" });
-    }
-}
-export const listMessagesByThreadId = query({
+const cloneMessageArgs = {
+    sourceThreadId: v.id("threads"),
+    targetThreadId: v.id("threads"),
+    // defaults to false, so searching for a message by userId will not find
+    // these copies
+    copyUserIdForVectorSearch: v.optional(v.boolean()),
+    // defaults to false, so tool calls & responses will be copied
+    excludeToolMessages: v.optional(v.boolean()),
+    // defaults to copying all messages, but you could just copy success messages.
+    statuses: v.optional(v.array(vMessageStatus)),
+    // stop at this message id
+    upToAndIncludingMessageId: v.optional(v.id("messages")),
+    // defaults to 0. the messages will be inserted starting at this order.
+    insertAtOrder: v.optional(v.number()),
+};
+export const cloneMessageBatch = internalMutation({
     args: {
-        threadId: v.id("threads"),
-        excludeToolMessages: v.optional(v.boolean()),
-        /** What order to sort the messages in. To get the latest, use "desc". */
-        order: v.union(v.literal("asc"), v.literal("desc")),
-        paginationOpts: v.optional(paginationOptsValidator),
-        statuses: v.optional(v.array(vMessageStatus)),
-        upToAndIncludingMessageId: v.optional(v.id("messages")),
+        ...cloneMessageArgs,
+        paginationOpts: paginationOptsValidator,
     },
     handler: async (ctx, args) => {
-        const statuses = args.statuses ?? vMessageStatus.members.map((m) => m.value);
-        const last = args.upToAndIncludingMessageId &&
-            (await ctx.db.get(args.upToAndIncludingMessageId));
-        assert(!last || last.threadId === args.threadId, "upToAndIncludingMessageId must be a message in the thread");
-        const toolOptions = args.excludeToolMessages ? [false] : [true, false];
-        const order = args.order ?? "desc";
-        const streams = toolOptions.flatMap((tool) => statuses.map((status) => stream(ctx.db, schema)
-            .query("messages")
-            .withIndex("threadId_status_tool_order_stepOrder", (q) => {
-            const qq = q
-                .eq("threadId", args.threadId)
+        const orderOffset = args.insertAtOrder ?? 0;
+        const result = await listMessagesByThreadIdHandler(ctx, {
+            threadId: args.sourceThreadId,
+            excludeToolMessages: args.excludeToolMessages,
+            order: "desc",
+            paginationOpts: args.paginationOpts,
+            statuses: args.statuses,
+            upToAndIncludingMessageId: args.upToAndIncludingMessageId,
+        });
+        const existing = result.page.length === 0
+            ? []
+            : await mergedStream([true, false].flatMap((tool) => messageStatuses.map((status) => stream(ctx.db, schema)
+                .query("messages")
+                .withIndex("threadId_status_tool_order_stepOrder", (q) => q
+                .eq("threadId", args.targetThreadId)
                 .eq("status", status)
-                .eq("tool", tool);
-            if (last) {
-                return qq.lte("order", last.order);
+                .eq("tool", tool)
+                .gte("order", result.page[0].order)
+                .lte("order", result.page[result.page.length - 1].order)))), ["order", "stepOrder"]).collect();
+        await Promise.all(result.page
+            .filter((m) => !existing.some((e) => e.order === m.order && e.stepOrder === m.stepOrder))
+            .map(async (m) => {
+            // update file refs
+            if (m.fileIds) {
+                await changeRefcount(ctx, [], m.fileIds);
             }
-            return qq;
-        })
-            .order(order)
-            .filterWith(
-        // We allow all messages on the same order.
-        async (m) => !last || m.order < last.order || m.order === last.order)));
-        const messages = await mergedStream(streams, [
-            "order",
-            "stepOrder",
-        ]).paginate(args.paginationOpts ?? {
-            numItems: DEFAULT_RECENT_MESSAGES,
-            cursor: null,
-        });
+            let embeddingId = undefined;
+            if (m.embeddingId) {
+                const vector = await ctx.db.get(m.embeddingId);
+                assert(vector, `Vector ${m.embeddingId} not found`);
+                const dimension = vector.vector.length;
+                validateVectorDimension(dimension);
+                embeddingId = await insertVector(ctx, dimension, {
+                    ...pick(vector, ["model", "table", "vector"]),
+                    userId: args.copyUserIdForVectorSearch
+                        ? vector.userId
+                        : undefined,
+                    threadId: args.targetThreadId,
+                });
+            }
+            await ctx.db.insert("messages", {
+                ...omit(m, [
+                    "_id",
+                    "_creationTime",
+                    "threadId",
+                    "order",
+                    "embeddingId",
+                ]),
+                embeddingId,
+                threadId: args.targetThreadId,
+                order: orderOffset + m.order,
+            });
+        }));
+        return {
+            numCopied: result.page.length,
+            continueCursor: result.continueCursor,
+            isDone: result.isDone,
+        };
+    },
+});
+export const cloneThread = action({
+    args: {
+        ...cloneMessageArgs,
+        batchSize: v.optional(v.number()),
+        // how many messages to copy
+        limit: v.optional(v.number()),
+    },
+    returns: v.number(),
+    handler: async (ctx, args) => {
+        let cursor = null;
+        let copiedSoFar = 0;
+        while (copiedSoFar < (args.limit ?? Infinity)) {
+            const numToCopy = Math.min(args.batchSize ?? DEFAULT_RECENT_MESSAGES, args.limit ?? Infinity - copiedSoFar);
+            const result = await ctx.runMutation(internal.messages.cloneMessageBatch, {
+                ...args,
+                paginationOpts: {
+                    cursor,
+                    numItems: numToCopy,
+                },
+            });
+            copiedSoFar += result.numCopied;
+            cursor = result.continueCursor;
+            if (result.isDone) {
+                break;
+            }
+        }
+        return copiedSoFar;
+    },
+});
+export const listMessagesByThreadIdArgs = {
+    threadId: v.id("threads"),
+    excludeToolMessages: v.optional(v.boolean()),
+    /** What order to sort the messages in. To get the latest, use "desc". */
+    order: v.union(v.literal("asc"), v.literal("desc")),
+    paginationOpts: v.optional(paginationOptsValidator),
+    statuses: v.optional(v.array(vMessageStatus)),
+    upToAndIncludingMessageId: v.optional(v.id("messages")),
+};
+export const listMessagesByThreadId = query({
+    args: listMessagesByThreadIdArgs,
+    handler: async (ctx, args) => {
+        const messages = await listMessagesByThreadIdHandler(ctx, args);
         return { ...messages, page: messages.page.map(publicMessage) };
     },
     returns: vPaginationResult(vMessageDoc),
 });
+async function listMessagesByThreadIdHandler(ctx, args) {
+    const statuses = args.statuses ?? vMessageStatus.members.map((m) => m.value);
+    const last = args.upToAndIncludingMessageId &&
+        (await ctx.db.get(args.upToAndIncludingMessageId));
+    assert(!last || last.threadId === args.threadId, "upToAndIncludingMessageId must be a message in the thread");
+    const toolOptions = args.excludeToolMessages ? [false] : [true, false];
+    const order = args.order ?? "desc";
+    const streams = toolOptions.flatMap((tool) => statuses.map((status) => stream(ctx.db, schema)
+        .query("messages")
+        .withIndex("threadId_status_tool_order_stepOrder", (q) => {
+        const qq = q
+            .eq("threadId", args.threadId)
+            .eq("status", status)
+            .eq("tool", tool);
+        if (last) {
+            return qq.lte("order", last.order);
+        }
+        return qq;
+    })
+        .order(order)
+        .filterWith(
+    // We allow all messages on the same order.
+    async (m) => !last || m.order <= last.order)));
+    const messages = await mergedStream(streams, ["order", "stepOrder"]).paginate(args.paginationOpts ?? {
+        numItems: DEFAULT_RECENT_MESSAGES,
+        cursor: null,
+    });
+    if (messages.page.length === 0) {
+        messages.isDone = true;
+    }
+    return messages;
+}
 export const getMessagesByIds = query({
-    args: {
-        messageIds: v.array(v.id("messages")),
-    },
+    args: { messageIds: v.array(v.id("messages")) },
     handler: async (ctx, args) => {
-        return await Promise.all(args.messageIds.map((id) => ctx.db.get(id)));
+        return (await Promise.all(args.messageIds.map((id) => ctx.db.get(id)))).map((m) => (m ? publicMessage(m) : null));
     },
     returns: v.array(v.union(v.null(), vMessageDoc)),
 });
@@ -320,10 +511,12 @@ export const searchMessages = action({
     args: {
         threadId: v.optional(v.id("threads")),
         searchAllMessagesForUserId: v.optional(v.string()),
-        beforeMessageId: v.optional(v.id("messages")),
+        targetMessageId: v.optional(v.id("messages")),
         embedding: v.optional(v.array(v.number())),
         embeddingModel: v.optional(v.string()),
         text: v.optional(v.string()),
+        textSearch: v.optional(v.boolean()),
+        vectorSearch: v.optional(v.boolean()),
         limit: v.number(),
         vectorScoreThreshold: v.optional(v.number()),
         messageRange: v.optional(v.object({ before: v.number(), after: v.number() })),
@@ -333,23 +526,34 @@ export const searchMessages = action({
         assert(args.searchAllMessagesForUserId || args.threadId, "Specify userId or threadId");
         const limit = args.limit;
         let textSearchMessages;
-        if (args.text) {
+        if (args.textSearch) {
             textSearchMessages = await ctx.runQuery(api.messages.textSearch, {
                 searchAllMessagesForUserId: args.searchAllMessagesForUserId,
                 threadId: args.threadId,
+                targetMessageId: args.targetMessageId,
                 text: args.text,
                 limit,
-                beforeMessageId: args.beforeMessageId,
             });
         }
-        if (args.embedding) {
-            const dimension = args.embedding.length;
-            if (!VectorDimensions.includes(dimension)) {
-                throw new Error(`Unsupported embedding dimension: ${dimension}`);
+        if (args.vectorSearch) {
+            let embedding = args.embedding;
+            let model = args.embeddingModel;
+            if (!embedding) {
+                if (args.targetMessageId) {
+                    const target = await ctx.runQuery(api.messages.getMessageSearchFields, {
+                        messageId: args.targetMessageId,
+                    });
+                    assert(target, "Target message embedding not found.");
+                    embedding = target.embedding;
+                    model = target.embeddingModel;
+                }
             }
-            const vectors = (await searchVectors(ctx, args.embedding, {
+            assert(embedding && model, "Embedding missing");
+            const dimension = embedding.length;
+            validateVectorDimension(dimension);
+            const vectors = (await searchVectors(ctx, embedding, {
                 dimension,
-                model: args.embeddingModel ?? "unknown",
+                model,
                 table: "messages",
                 searchAllMessagesForUserId: args.searchAllMessagesForUserId,
                 threadId: args.threadId,
@@ -372,7 +576,7 @@ export const searchMessages = action({
                 embeddingIds,
                 textSearchMessages: textSearchMessages?.filter((m) => !embeddingIds.includes(m.embeddingId)),
                 messageRange: args.messageRange ?? DEFAULT_MESSAGE_RANGE,
-                beforeMessageId: args.beforeMessageId,
+                beforeMessageId: args.targetMessageId,
                 limit,
             });
             return messages;
@@ -414,11 +618,11 @@ export const _fetchSearchMessages = internalQuery({
             .map(publicMessage);
         messages.push(...(args.textSearchMessages ?? []));
         // TODO: prioritize more recent messages
-        messages.sort((a, b) => a.order - b.order);
+        messages = sorted(messages);
         messages = messages.slice(0, args.limit);
         // Fetch the surrounding messages
         if (!threadId) {
-            return messages.sort((a, b) => a.order - b.order);
+            return messages;
         }
         const included = {};
         for (const m of messages) {
@@ -469,7 +673,7 @@ export const _fetchSearchMessages = internalQuery({
                 messages.push(publicMessage(r));
             }
         }
-        return messages.sort((a, b) => a.order - b.order);
+        return sorted(messages);
     },
 });
 // returns ranges of messages in order of text search relevance,
@@ -478,21 +682,24 @@ export const textSearch = query({
     args: {
         threadId: v.optional(v.id("threads")),
         searchAllMessagesForUserId: v.optional(v.string()),
-        text: v.string(),
+        text: v.optional(v.string()),
+        targetMessageId: v.optional(v.id("messages")),
         limit: v.number(),
-        beforeMessageId: v.optional(v.id("messages")),
     },
     handler: async (ctx, args) => {
         assert(args.searchAllMessagesForUserId || args.threadId, "Specify userId or threadId");
-        const beforeMessage = args.beforeMessageId && (await ctx.db.get(args.beforeMessageId));
-        const order = beforeMessage?.order;
+        const targetMessage = args.targetMessageId && (await ctx.db.get(args.targetMessageId));
+        const order = targetMessage?.order;
+        const text = args.text || targetMessage?.text;
+        if (!text) {
+            console.warn("No text to search", targetMessage, args.text);
+            return [];
+        }
         const messages = await ctx.db
             .query("messages")
             .withSearchIndex("text_search", (q) => args.searchAllMessagesForUserId
-            ? q
-                .search("text", args.text)
-                .eq("userId", args.searchAllMessagesForUserId)
-            : q.search("text", args.text).eq("threadId", args.threadId))
+            ? q.search("text", text).eq("userId", args.searchAllMessagesForUserId)
+            : q.search("text", text).eq("threadId", args.threadId))
             // Just in case tool messages slip through
             .filter((q) => {
             const qq = q.eq(q.field("tool"), false);
@@ -503,12 +710,38 @@ export const textSearch = query({
         })
             .take(args.limit);
         return messages
-            .filter((m) => !beforeMessage ||
-            m.order < beforeMessage.order ||
-            (m.order === beforeMessage.order &&
-                m.stepOrder < beforeMessage.stepOrder))
+            .filter((m) => !targetMessage ||
+            m.order < targetMessage.order ||
+            (m.order === targetMessage.order &&
+                m.stepOrder < targetMessage.stepOrder))
             .map(publicMessage);
     },
     returns: v.array(vMessageDoc),
 });
+export const getMessageSearchFields = query({
+    args: {
+        messageId: v.id("messages"),
+    },
+    returns: v.object({
+        text: v.optional(v.string()),
+        embedding: v.optional(v.array(v.number())),
+        embeddingModel: v.optional(v.string()),
+    }),
+    handler: async (ctx, args) => {
+        const message = await ctx.db.get(args.messageId);
+        const text = message?.text;
+        let embedding = undefined;
+        let embeddingModel = undefined;
+        if (message?.embeddingId) {
+            const target = await ctx.db.get(message.embeddingId);
+            embedding = target?.vector;
+            embeddingModel = target?.model;
+        }
+        return {
+            text,
+            embedding,
+            embeddingModel,
+        };
+    },
+});
 //# sourceMappingURL=messages.js.map