feat: patch for assistant ui #3472

package.json

@@ -116,6 +116,7 @@
     "geist": "^1.7.0",
     "gsap": "^3.14.2",
     "input-otp": "1.4.2",
+    "jsonrepair": "^3.13.2",
     "katex": "^0.16.33",
     "lucide-react": "^0.575.0",
     "mathlive": "^0.108.2",

src/components/assistant-ui/thread.tsx

@@ -1240,36 +1240,29 @@ const AssistantMessage: FC = () => {
   );
 };
 
-const CONTINUE_RESPONSE_CHAR_THRESHOLD = 150;
-
-/** Sentence-ending punctuation that indicates an LLM likely finished naturally. */
-const ENDS_WITH_SENTENCE_PUNCTUATION = /[.!?]$/;
-
-/** Trailing citation tags to strip before punctuation check (e.g. <citation>Source</citation>). */
-const TRAILING_CITATION_TAGS = /(?:\s*<citation>[\s\S]*?<\/citation>\s*)+$/;
+/**
+ * Max characters above which we never show the Continue button.
+ * Long responses (citations, tables, mermaid, multi-paragraph) are assumed complete.
+ * We avoid punctuation-based heuristics since the agent often ends with
+ * <citation>, tables, mermaid blocks, etc.
+ */
+const CONTINUE_MAX_CHARS = 600;
 
 const AssistantActionBar: FC = () => {
   const { createCard, isCreating } = useCreateCardFromMessage({ debounceMs: 300 });
   const message = useMessage();
   const api = useAssistantApi();
 
-  const { textLength, textContent, endsWithPunctuation } = useMemo(() => {
+  const { textLength, textContent } = useMemo(() => {
     const textParts = message.content.filter(
       (part): part is { type: "text"; text: string } => part.type === "text"
     );
     const length = textParts.reduce((sum, part) => sum + (part.text?.length ?? 0), 0);
     const content = textParts.map((part) => part.text ?? "").join("\n\n");
-    const trimmed = content.trim();
-    const withoutTrailingCitations = trimmed.replace(TRAILING_CITATION_TAGS, "").trim();
-    const endsWithPunctuation =
-      withoutTrailingCitations.length > 0 &&
-      ENDS_WITH_SENTENCE_PUNCTUATION.test(withoutTrailingCitations);
-    return { textLength: length, textContent: content, endsWithPunctuation };
+    return { textLength: length, textContent: content };
   }, [message.content]);
 
-  const showContinueButton =
-    textLength > 0 &&
-    (textLength < CONTINUE_RESPONSE_CHAR_THRESHOLD || !endsWithPunctuation);
+  const showContinueButton = textLength > 0 && textLength < CONTINUE_MAX_CHARS;
 
   const handleContinueClick = useCallback(() => {
     api?.thread().append({

src/lib/ai/tools/__tests__/edit-item-tool.test.ts

@@ -0,0 +1,148 @@
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+
+const mockWorkspaceWorker = vi.fn();
+const mockLoadStateForTool = vi.fn();
+const mockResolveItem = vi.fn();
+const mockGetVirtualPath = vi.fn();
+
+vi.mock("@/lib/ai/workers", () => ({
+  workspaceWorker: (...args: unknown[]) => mockWorkspaceWorker(...args),
+}));
+
+vi.mock("@/lib/ai/tools/tool-utils", () => ({
+  loadStateForTool: (...args: unknown[]) => mockLoadStateForTool(...args),
+  resolveItem: (...args: unknown[]) => mockResolveItem(...args),
+}));
+
+vi.mock("@/lib/utils/virtual-workspace-fs", () => ({
+  getVirtualPath: (...args: unknown[]) => mockGetVirtualPath(...args),
+}));
+
+describe("createEditItemTool", () => {
+  let createEditItemTool: (ctx: {
+    workspaceId: string;
+    userId: string;
+    activeFolderId?: string;
+    threadId?: string | null;
+  }) => any;
+
+  const ctx = {
+    workspaceId: "ws-1",
+    userId: "user-1",
+    activeFolderId: undefined,
+    threadId: "thread-1",
+  };
+
+  const noteItem = {
+    id: "note-1",
+    type: "note",
+    name: "My Note",
+    data: {},
+  };
+
+  beforeAll(async () => {
+    const mod = await import("@/lib/ai/tools/edit-item-tool");
+    createEditItemTool = mod.createEditItemTool;
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockLoadStateForTool.mockResolvedValue({ success: true, state: { items: [noteItem] } });
+    mockResolveItem.mockReturnValue(noteItem);
+    mockWorkspaceWorker.mockResolvedValue({ success: true, itemId: "note-1", message: "ok" });
+    mockGetVirtualPath.mockReturnValue("notes/My Note.md");
+  });
+
+  it("fails fast when itemName is missing", async () => {
+    const tool: any = createEditItemTool(ctx);
+    const result = await tool.execute({ itemName: "", oldString: "a", newString: "b" });
+    expect(result.success).toBe(false);
+    expect(result.message).toMatch(/itemName is required/i);
+  });
+
+  it("returns loadStateForTool failure", async () => {
+    mockLoadStateForTool.mockResolvedValueOnce({ success: false, message: "No workspace context available" });
+    const tool: any = createEditItemTool(ctx);
+    const result = await tool.execute({ itemName: "My Note", oldString: "a", newString: "b" });
+    expect(result.success).toBe(false);
+    expect(result.message).toMatch(/No workspace context available/);
+    expect(mockWorkspaceWorker).not.toHaveBeenCalled();
+  });
+
+  it("returns item not found message when resolveItem misses", async () => {
+    mockResolveItem.mockReturnValueOnce(undefined);
+    const tool: any = createEditItemTool(ctx);
+    const result = await tool.execute({ itemName: "Unknown", oldString: "a", newString: "b" });
+    expect(result.success).toBe(false);
+    expect(result.message).toMatch(/Could not find item "Unknown"/);
+  });
+
+  it("requires disambiguation when multiple exact-name candidates exist", async () => {
+    const duplicate = { ...noteItem, id: "note-2" };
+    mockLoadStateForTool.mockResolvedValueOnce({
+      success: true,
+      state: { items: [noteItem, duplicate] },
+    });
+    mockResolveItem.mockReturnValueOnce(noteItem);
+    mockGetVirtualPath
+      .mockReturnValueOnce("notes/My Note.md")
+      .mockReturnValueOnce("folder/My Note.md");
+
+    const tool: any = createEditItemTool(ctx);
+    const result = await tool.execute({ itemName: "My Note", oldString: "a", newString: "b" });
+    expect(result.success).toBe(false);
+    expect(result.message).toMatch(/Multiple items named/);
+    expect(result.message).toMatch(/Disambiguate using path/);
+  });
+
+  it("rejects non-editable item types", async () => {
+    const imageItem = { id: "img-1", type: "image", name: "Figure", data: {} };
+    mockResolveItem.mockReturnValueOnce(imageItem);
+    mockLoadStateForTool.mockResolvedValueOnce({ success: true, state: { items: [imageItem] } });
+
+    const tool: any = createEditItemTool(ctx);
+    const result = await tool.execute({ itemName: "Figure", oldString: "a", newString: "b" });
+    expect(result.success).toBe(false);
+    expect(result.message).toMatch(/not editable/);
+  });
+
+  it("calls workspaceWorker edit and returns renamed itemName", async () => {
+    const tool: any = createEditItemTool(ctx);
+    const result = await tool.execute({
+      itemName: "My Note",
+      oldString: "old",
+      newString: "new",
+      newName: "Renamed Note",
+      replaceAll: true,
+      sources: [{ title: "Ref", url: "https://example.com" }],
+    });
+
+    expect(mockWorkspaceWorker).toHaveBeenCalledWith(
+      "edit",
+      expect.objectContaining({
+        workspaceId: "ws-1",
+        itemId: "note-1",
+        itemType: "note",
+        oldString: "old",
+        newString: "new",
+        replaceAll: true,
+        newName: "Renamed Note",
+      })
+    );
+    expect(result.success).toBe(true);
+    expect(result.itemName).toBe("Renamed Note");
+  });
+
+  it("passes through worker failures", async () => {
+    mockWorkspaceWorker.mockResolvedValueOnce({
+      success: false,
+      message: "Found multiple matches for oldString",
+    });
+    const tool: any = createEditItemTool(ctx);
+    const result = await tool.execute({ itemName: "My Note", oldString: "a", newString: "b" });
+
+    expect(result.success).toBe(false);
+    expect(result.message).toMatch(/multiple matches/i);
+  });
+});
+

src/lib/ai/tools/read-workspace.ts

@@ -4,8 +4,10 @@ import { loadStateForTool } from "./tool-utils";
 import { fuzzyMatchItem } from "./tool-utils";
 import { resolveItemByPath } from "./workspace-search-utils";
 import { formatItemContent } from "@/lib/utils/format-workspace-context";
+import { getNoteContentAsMarkdown } from "@/lib/utils/format-workspace-context";
 import { getVirtualPath } from "@/lib/utils/virtual-workspace-fs";
 import type { WorkspaceToolContext } from "./workspace-tools";
+import type { NoteData } from "@/lib/workspace-state/types";
 
 const DEFAULT_LIMIT = 500;
 const MAX_LIMIT = 2000;
@@ -122,7 +124,13 @@ export function createReadWorkspaceTool(ctx: WorkspaceToolContext) {
                     ? { pageStart, pageEnd }
                     : undefined;
 
-            const fullContent = formatItemContent(item, pdfPageRange);
+            // Keep note reads aligned with editItem matching source.
+            // editItem for notes matches against serialized markdown content only
+            // (without appended metadata blocks like Sources).
+            const fullContent =
+                item.type === "note"
+                    ? getNoteContentAsMarkdown(item.data as NoteData)
+                    : formatItemContent(item, pdfPageRange);
             const allLines = fullContent.split(/\r?\n/);
             const totalLines = allLines.length;
             const startIdx = Math.max(0, lineStart - 1);