diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 0edf432d..d63ef489 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -137,6 +137,7 @@ export type ComponentApi = "mutation", "internal", { + abandonIfPromptMissing?: boolean; agentName?: string; embeddings?: { dimension: diff --git a/src/component/messages.test.ts b/src/component/messages.test.ts index be00786f..ce380eb8 100644 --- a/src/component/messages.test.ts +++ b/src/component/messages.test.ts @@ -1328,3 +1328,96 @@ describe("late saves racing a failed pending message (issue #320)", () => { expect(stream?.state.kind).toBe("aborted"); }); }); + +describe("deleting a message aborts generation writing to it (issue #300)", () => { + test("a stream at the deleted order is aborted", async () => { + const t = initConvexTest(); + const thread = await t.mutation(api.threads.createThread, { userId: "u" }); + const threadId = thread._id as Id<"threads">; + + const { messages } = await t.mutation(api.messages.addMessages, { + threadId, + messages: [{ message: { role: "user", content: "hello" } }], + }); + const prompt = messages[0]; + + await t.mutation(api.streams.create, { + threadId, + order: prompt.order, + stepOrder: prompt.stepOrder + 1, + userId: "u", + agentName: "a", + model: "m", + provider: "p", + format: "UIMessageChunk", + }); + + await t.mutation(api.messages.deleteByIds, { + messageIds: [prompt._id as Id<"messages">], + }); + + const streaming = await t.query(api.streams.list, { + threadId, + statuses: ["streaming"], + }); + const aborted = await t.query(api.streams.list, { + threadId, + statuses: ["aborted"], + }); + expect(streaming).toHaveLength(0); + expect(aborted).toHaveLength(1); + }); +}); + +describe("abandoning a save whose prompt was deleted (issue #300)", () => { + test("abandons instead of throwing when the caller opts in", async () => { + const t = initConvexTest(); + const thread = await t.mutation(api.threads.createThread, { userId: "u" }); + const threadId = thread._id as Id<"threads">; + + const { messages } = await t.mutation(api.messages.addMessages, { + threadId, + messages: [{ message: { role: "user", content: "hello" } }], + }); + const promptMessageId = messages[0]._id as Id<"messages">; + + await t.mutation(api.messages.deleteByIds, { messageIds: [promptMessageId] }); + + const saved = await t.mutation(api.messages.addMessages, { + threadId, + promptMessageId, + abandonIfPromptMissing: true, + messages: [{ message: { role: "assistant", content: "answer" } }], + }); + expect(saved.messages).toEqual([]); + + // Nothing was grafted onto the thread. + const all = await t.query(api.messages.listMessagesByThreadId, { + threadId, + order: "asc", + paginationOpts: { cursor: null, numItems: 10 }, + }); + expect(all.page).toHaveLength(0); + }); + + test("still throws for a caller that did not opt in", async () => { + const t = initConvexTest(); + const thread = await t.mutation(api.threads.createThread, { userId: "u" }); + const threadId = thread._id as Id<"threads">; + + const { messages } = await t.mutation(api.messages.addMessages, { + threadId, + messages: [{ message: { role: "user", content: "hello" } }], + }); + const promptMessageId = messages[0]._id as Id<"messages">; + await t.mutation(api.messages.deleteByIds, { messageIds: [promptMessageId] }); + + await expect( + t.mutation(api.messages.addMessages, { + threadId, + promptMessageId, + messages: [{ message: { role: "assistant", content: "answer" } }], + }), + ).rejects.toThrow("not found"); + }); +}); diff --git a/src/component/messages.ts b/src/component/messages.ts index e1e0daa0..9b87d257 100644 --- a/src/component/messages.ts +++ b/src/component/messages.ts @@ -44,6 +44,7 @@ import { getStreamingMessagesWithMetadata, finishHandler, releaseStreamFileOwnershipByIds, + abortStreamsAtOrder, } from "./streams.js"; import { partial } from "convex-helpers/validators"; @@ -65,21 +66,45 @@ export async function deleteMessage( } } +/** + * Deleting a message strands any generation still writing to its order, which + * would otherwise only surface as a missing-parent failure when that generation + * finalizes. Aborting the stream lets the in-flight run stop on its own. + */ +async function abortStreamsForDeleted( + ctx: MutationCtx, + deleted: (Doc<"messages"> | null)[], +) { + const seen = new Set(); + for (const message of deleted) { + if (!message) continue; + const key = `${message.threadId}:${message.order}`; + if (seen.has(key)) continue; + seen.add(key); + await abortStreamsAtOrder(ctx, { + threadId: message.threadId, + order: message.order, + reason: "Message deleted", + }); + } +} + export const deleteByIds = mutation({ args: { messageIds: v.array(v.id("messages")) }, returns: v.array(v.id("messages")), handler: async (ctx, args) => { - const deletedMessageIds = await Promise.all( + const deleted = await Promise.all( args.messageIds.map(async (id) => { const message = await ctx.db.get("messages", id); if (message) { await deleteMessage(ctx, message); - return id; + return message; } return null; }), ); - return deletedMessageIds.filter((id) => id !== null); + await abortStreamsForDeleted(ctx, deleted); + return deleted.filter((m) => m !== null).map((m) => m._id); }, }); @@ -145,6 +170,7 @@ export const deleteByOrder = mutation({ }) .take(64); await Promise.all(messages.map((m) => deleteMessage(ctx, m))); + await abortStreamsForDeleted(ctx, messages); return { isDone: messages.length < 64, lastOrder: messages.at(-1)?.order, @@ -157,6 +183,12 @@ const addMessagesArgs = { userId: v.optional(v.string()), threadId: v.id("threads"), promptMessageId: v.optional(v.id("messages")), + /** + * For saves that belong to a run anchored on promptMessageId: if that + * message is gone the run is obsolete, so abandon the save instead of + * throwing. A caller passing an id that never existed still gets an error. + */ + abandonIfPromptMissing: v.optional(v.boolean()), order: v.optional(v.union(v.number(), v.literal("next"))), agentName: v.optional(v.string()), messages: v.array(vMessageWithMetadataInternal), @@ -203,6 +235,7 @@ async function addMessagesHandler( finishStreamId, messages, promptMessageId, + abandonIfPromptMissing, order: requestedOrder, pendingMessageId, hideFromUserIdSearch, @@ -267,6 +300,9 @@ async function addMessagesHandler( const maxMessage = await getMaxMessage(ctx, threadId, order); stepOrder = maxMessage?.stepOrder ?? -1; } else if (promptMessageId) { + if (!promptMessage && abandonIfPromptMissing) { + return { messages: [] }; + } assert(promptMessage, `Parent message ${promptMessageId} not found`); if (promptMessage.status === "failed") { fail = true; diff --git a/src/component/streams.ts b/src/component/streams.ts index afa49302..2c96c4b7 100644 --- a/src/component/streams.ts +++ b/src/component/streams.ts @@ -197,24 +197,29 @@ function publicStreamMessage(m: Doc<"streamingMessages">): StreamMessage { }; } +export async function abortStreamsAtOrder( + ctx: MutationCtx, + args: { threadId: Id<"threads">; order: number; reason: string }, +) { + const streams = await ctx.db + .query("streamingMessages") + .withIndex("threadId_state_order_stepOrder", (q) => + q + .eq("threadId", args.threadId) + .eq("state.kind", "streaming") + .eq("order", args.order), + ) + .take(100); + for (const stream of streams) { + await abortById(ctx, { streamId: stream._id, reason: args.reason }); + } + return streams.length > 0; +} + export const abortByOrder = mutation({ args: { threadId: v.id("threads"), order: v.number(), reason: v.string() }, returns: v.boolean(), - handler: async (ctx, args) => { - const streams = await ctx.db - .query("streamingMessages") - .withIndex("threadId_state_order_stepOrder", (q) => - q - .eq("threadId", args.threadId) - .eq("state.kind", "streaming") - .eq("order", args.order), - ) - .take(100); - for (const stream of streams) { - await abortById(ctx, { streamId: stream._id, reason: args.reason }); - } - return streams.length > 0; - }, + handler: abortStreamsAtOrder, }); export const abort = mutation({ diff --git a/src/vercel/client/start.ts b/src/vercel/client/start.ts index f0af0e45..4d7c8642 100644 --- a/src/vercel/client/start.ts +++ b/src/vercel/client/start.ts @@ -386,6 +386,7 @@ export async function startGeneration< threadId, agentName: opts.agentName, promptMessageId, + abandonIfPromptMissing: true, pendingMessageId, messages: serialized.messages, embeddings,