Fix three UI / streaming / search bugs

package.json

@@ -112,7 +112,7 @@
     "clsx": "2.1.1",
     "convex": "1.35.1",
     "convex-helpers": "0.1.114",
-    "convex-test": "0.0.44",
+    "convex-test": "0.0.51",
     "dayjs": "1.11.20",
     "dotenv": "16.6.1",
     "eslint": "9.39.4",

src/UIMessages.ts

@@ -337,7 +337,7 @@ function createUserUIMessage<
 
   const parts: UIMessage<METADATA, DATA_PARTS, TOOLS>["parts"] = [];
   if (text && !nonStringContent.length) {
-    parts.push({ type: "text", text });
+    parts.push({ type: "text", text, ...partCommon });
   }
   for (const contentPart of nonStringContent) {
     switch (contentPart.type) {

src/component/messages.test.ts

@@ -623,6 +623,65 @@ describe("agent", () => {
     expect(result.lastStepOrder).toBe(1);
   });
 
+  test("textSearch returns cross-thread matches even when target order is lower (#256)", async () => {
+    const t = initConvexTest();
+    const userId = "search-user-1";
+
+    // Old thread: matching message lives at order 5.
+    const threadA = await t.mutation(api.threads.createThread, { userId });
+    for (let i = 0; i < 5; i++) {
+      await t.mutation(api.messages.addMessages, {
+        threadId: threadA._id as Id<"threads">,
+        messages: [{ message: { role: "user", content: `padding ${i}` } }],
+      });
+    }
+    await t.mutation(api.messages.addMessages, {
+      threadId: threadA._id as Id<"threads">,
+      messages: [
+        { message: { role: "user", content: "high-ticket coaches are great" } },
+      ],
+    });
+
+    // New thread: target lives at order 2 — non-zero so the buggy DB-level
+    // `lte(order)` filter actually engages and drops threadA's order-5 match.
+    const threadB = await t.mutation(api.threads.createThread, { userId });
+    await t.mutation(api.messages.addMessages, {
+      threadId: threadB._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "intro one" } }],
+    });
+    await t.mutation(api.messages.addMessages, {
+      threadId: threadB._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "intro two" } }],
+    });
+    const { messages: targetMsgs } = await t.mutation(
+      api.messages.addMessages,
+      {
+        threadId: threadB._id as Id<"threads">,
+        messages: [
+          { message: { role: "user", content: "high-ticket coaches" } },
+        ],
+      },
+    );
+    const targetMessageId = targetMsgs[0]._id as Id<"messages">;
+
+    const results = await t.query(api.messages.textSearch, {
+      searchAllMessagesForUserId: userId,
+      text: "high-ticket coaches",
+      targetMessageId,
+      limit: 10,
+    });
+
+    // Pre-fix, threadA's match (order 5) was dropped because the DB-level
+    // `lte(order, 2)` filter was applied to all threads using threadB's
+    // target order. Post-fix, cross-thread results bypass that filter.
+    const fromThreadA = results.filter((m) => m.threadId === threadA._id);
+    expect(fromThreadA.length).toBeGreaterThan(0);
+
+    // The target message itself must still be excluded from its own thread's results.
+    const targetInResults = results.find((m) => m._id === targetMessageId);
+    expect(targetInResults).toBeUndefined();
+  });
+
   test("deleteByOrder handles empty result set", async () => {
     const t = convexTest(schema, modules);
     const thread = await t.mutation(api.threads.createThread, {

src/component/messages.ts

@@ -821,16 +821,18 @@ export const _fetchSearchMessages = internalQuery({
         ),
       )
     )
-      .filter(
-        (m): m is Doc<"messages"> =>
-          m !== undefined &&
-          m !== null &&
-          !m.tool &&
-          (!beforeMessage ||
-            m.order < beforeMessage.order ||
-            (m.order === beforeMessage.order &&
-              m.stepOrder < beforeMessage.stepOrder)),
-      )
+      .filter((m): m is Doc<"messages"> => {
+        if (m === undefined || m === null || m.tool) return false;
+        if (!beforeMessage) return true;
+        // Order is only comparable within a single thread; cross-thread
+        // results have independent order sequences and must pass through.
+        if (m.threadId !== beforeMessage.threadId) return true;
+        return (
+          m.order < beforeMessage.order ||
+          (m.order === beforeMessage.order &&
+            m.stepOrder < beforeMessage.stepOrder)
+        );
+      })
       .map(publicMessage);
     messages.push(...(args.textSearchMessages ?? []));
     // TODO: prioritize more recent messages
@@ -928,20 +930,31 @@ export const textSearch = query({
       // Just in case tool messages slip through
       .filter((q) => {
         const qq = q.eq(q.field("tool"), false);
-        if (order) {
+        // Order is only comparable within a single thread, so when searching
+        // across threads for a user we skip the DB-level order filter and
+        // apply it per-thread below.
+        if (order && !args.searchAllMessagesForUserId) {
           return q.and(qq, q.lte(q.field("order"), order));
         }
         return qq;
       })
       .take(args.limit);
+    // Tradeoff for cross-thread search: in-thread matches with order >=
+    // targetMessage.order are dropped here rather than at the DB layer, so
+    // when there are many cross-thread matches they can edge out valid
+    // older in-thread matches. Acceptable for now; revisit with a merged
+    // per-thread + cross-thread query if it becomes an issue in practice.
     return messages
-      .filter(
-        (m) =>
-          !targetMessage ||
+      .filter((m) => {
+        if (!targetMessage) return true;
+        // Different threads have independent order sequences; don't compare across them.
+        if (m.threadId !== targetMessage.threadId) return true;
+        return (
           m.order < targetMessage.order ||
           (m.order === targetMessage.order &&
-            m.stepOrder < targetMessage.stepOrder),
-      )
+            m.stepOrder < targetMessage.stepOrder)
+        );
+      })
       .map(publicMessage);
   },
   returns: v.array(vMessageDoc),

src/toUIMessages.test.ts

@@ -32,7 +32,11 @@ describe("toUIMessages", () => {
     expect(uiMessages).toHaveLength(1);
     expect(uiMessages[0].role).toBe("user");
     expect(uiMessages[0].text).toBe("Hello!");
-    expect(uiMessages[0].parts[0]).toEqual({ type: "text", text: "Hello!" });
+    expect(uiMessages[0].parts[0]).toEqual({
+      type: "text",
+      text: "Hello!",
+      state: "done",
+    });
   });
 
   it("handles assistant message", () => {