@tambo-ai/react 0.68.0 → 0.69.1

package/esm/providers/tambo-context-attachment-provider.test.js

@@ -1,29 +1,33 @@
 import { act, renderHook, waitFor } from "@testing-library/react";
 import React from "react";
 import { TamboContextAttachmentProvider, useTamboContextAttachment, } from "./tambo-context-attachment-provider";
-import { TamboContextHelpersProvider, useTamboContextHelpers, } from "./tambo-context-helpers-provider";
+import { useTamboContextHelpers } from "./tambo-context-helpers-provider";
+// Mock the context helpers provider
+jest.mock("./tambo-context-helpers-provider");
+const mockAddContextHelper = jest.fn();
+const mockRemoveContextHelper = jest.fn();
+const CONTEXT_ATTACHMENTS_HELPER_KEY = "contextAttachments";
+beforeEach(() => {
+    jest.clearAllMocks();
+    useTamboContextHelpers.mockReturnValue({
+        addContextHelper: mockAddContextHelper,
+        removeContextHelper: mockRemoveContextHelper,
+    });
+});
 /**
  * Test suite for TamboContextAttachmentProvider
  *
  * Tests the context attachment feature which allows:
- * - Visual context badges above message input
- * - Automatic context helper registration/unregistration
- * - Custom suggestions that override auto-generated ones
- * - Dynamic context data customization via getContextHelperData
+ * - Adding context attachments that will be sent with the next message
+ * - Automatic registration/deregistration of context helpers
  */
 describe("TamboContextAttachmentProvider", () => {
-    beforeEach(() => {
-        jest.clearAllMocks();
-    });
     /**
-     * Base wrapper with both TamboContextAttachmentProvider and TamboContextHelpersProvider
-     * since context attachments need access to context helpers API
-     * @param getContextHelperData - Optional custom function to get context helper data
-     * @returns A React component that wraps children with the necessary providers
+     * Base wrapper with TamboContextAttachmentProvider
+     * @returns A React component that wraps children with the provider
      */
-    const createWrapper = (getContextHelperData) => {
-        const Wrapper = ({ children }) => (React.createElement(TamboContextHelpersProvider, null,
-            React.createElement(TamboContextAttachmentProvider, { getContextHelperData: getContextHelperData }, children)));
+    const createWrapper = () => {
+        const Wrapper = ({ children }) => (React.createElement(TamboContextAttachmentProvider, null, children));
         Wrapper.displayName = "TestWrapper";
         return Wrapper;
     };
@@ -39,8 +43,6 @@ describe("TamboContextAttachmentProvider", () => {
             expect(result.current).toHaveProperty("addContextAttachment");
             expect(result.current).toHaveProperty("removeContextAttachment");
             expect(result.current).toHaveProperty("clearContextAttachments");
-            expect(result.current).toHaveProperty("customSuggestions");
-            expect(result.current).toHaveProperty("setCustomSuggestions");
         });
         /**
          * Hook should throw error when used outside of provider
@@ -58,154 +60,318 @@ describe("TamboContextAttachmentProvider", () => {
     });
     describe("Adding Context Attachments", () => {
         /**
-         * Should add a context attachment with auto-generated ID
+         * Should add a context attachment and register/update the merged context helper
          */
-        it("should add a context attachment", () => {
+        it("should add a context attachment", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
+            let attachment;
             act(() => {
-                result.current.addContextAttachment({
-                    name: "Button.tsx",
-                    metadata: { filePath: "/src/Button.tsx" },
+                attachment = result.current.addContextAttachment({
+                    context: "selectedFile",
+                    displayName: "Button.tsx",
+                    type: "file",
                 });
             });
             expect(result.current.attachments).toHaveLength(1);
             expect(result.current.attachments[0]).toMatchObject({
-                name: "Button.tsx",
-                metadata: { filePath: "/src/Button.tsx" },
+                displayName: "Button.tsx",
+                context: "selectedFile",
+                type: "file",
             });
             expect(result.current.attachments[0].id).toBeDefined();
+            expect(attachment.id).toBe(result.current.attachments[0].id);
+            // Wait for useEffect to run and register the merged helper with the attachment
+            await waitFor(() => {
+                const lastCall = mockAddContextHelper.mock.calls
+                    .filter((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)
+                    .pop();
+                if (lastCall) {
+                    const helperFn = lastCall[1];
+                    const helperResult = helperFn();
+                    expect(helperResult).not.toBeNull();
+                    expect(helperResult).toHaveLength(1);
+                }
+                else {
+                    throw new Error("Helper not registered");
+                }
+            });
+            // Verify the helper function returns the merged attachments array
+            const helperFn = mockAddContextHelper.mock.calls
+                .filter((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)
+                .pop()?.[1];
+            expect(helperFn).toBeDefined();
+            const helperResult = helperFn();
+            expect(helperResult).toEqual([
+                {
+                    id: attachment.id,
+                    displayName: "Button.tsx",
+                    context: "selectedFile",
+                    type: "file",
+                },
+            ]);
         });
         /**
-         * Should add multiple different attachments
+         * Should add a context attachment without displayName
          */
-        it("should add multiple context attachments", () => {
+        it("should add a context attachment without displayName", () => {
+            const { result } = renderHook(() => useTamboContextAttachment(), {
+                wrapper: createWrapper(),
+            });
+            let attachment;
+            act(() => {
+                attachment = result.current.addContextAttachment({
+                    context: "selectedFile",
+                });
+            });
+            expect(result.current.attachments).toHaveLength(1);
+            expect(result.current.attachments[0]).toMatchObject({
+                context: "selectedFile",
+            });
+            expect(result.current.attachments[0].displayName).toBeUndefined();
+            expect(attachment.id).toBe(result.current.attachments[0].id);
+        });
+        /**
+         * Should add multiple different context attachments and update the merged helper
+         */
+        it("should add multiple context attachments", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
             act(() => {
                 result.current.addContextAttachment({
-                    name: "Button.tsx",
+                    context: "file1",
+                    displayName: "Button.tsx",
+                    type: "file",
                 });
                 result.current.addContextAttachment({
-                    name: "Card.tsx",
+                    context: "file2",
+                    displayName: "Card.tsx",
+                    type: "file",
                 });
             });
             expect(result.current.attachments).toHaveLength(2);
-            expect(result.current.attachments[0].name).toBe("Button.tsx");
-            expect(result.current.attachments[1].name).toBe("Card.tsx");
+            expect(result.current.attachments[0].displayName).toBe("Button.tsx");
+            expect(result.current.attachments[1].displayName).toBe("Card.tsx");
+            // Wait for useEffect to run and update the merged helper
+            await waitFor(() => {
+                // The helper should be called/updated for each attachment change
+                expect(mockAddContextHelper).toHaveBeenCalledWith(CONTEXT_ATTACHMENTS_HELPER_KEY, expect.any(Function));
+            });
+            // Verify the helper function returns all attachments
+            const helperFn = mockAddContextHelper.mock.calls
+                .filter((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)
+                .pop()?.[1];
+            expect(helperFn).toBeDefined();
+            const helperResult = helperFn();
+            expect(helperResult).toHaveLength(2);
+            expect(helperResult[0].displayName).toBe("Button.tsx");
+            expect(helperResult[1].displayName).toBe("Card.tsx");
         });
         /**
-         * Should prevent duplicates with the same name
+         * Should allow multiple attachments with the same context value
          */
-        it("should prevent duplicate attachments with same name", () => {
+        it("should allow multiple attachments with same context value", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
             act(() => {
                 result.current.addContextAttachment({
-                    name: "Button.tsx",
+                    context: "selectedFile",
+                    displayName: "Button.tsx",
+                    type: "file",
                 });
                 result.current.addContextAttachment({
-                    name: "Button.tsx",
+                    context: "selectedFile",
+                    displayName: "Card.tsx",
+                    type: "file",
                 });
             });
-            expect(result.current.attachments).toHaveLength(1);
+            expect(result.current.attachments).toHaveLength(2);
+            // Wait for useEffect to update the merged helper
+            await waitFor(() => {
+                expect(mockAddContextHelper).toHaveBeenCalledWith(CONTEXT_ATTACHMENTS_HELPER_KEY, expect.any(Function));
+            });
+            // Verify both attachments are included in the merged result
+            const helperFn = mockAddContextHelper.mock.calls
+                .filter((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)
+                .pop()?.[1];
+            expect(helperFn).toBeDefined();
+            const helperResult = helperFn();
+            expect(helperResult).toHaveLength(2);
         });
         /**
-         * Should support optional icon property
+         * Should support optional type property
          */
-        it("should support icon property", () => {
+        it("should support optional type property", () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
-            const icon = React.createElement("span", null, "\uD83D\uDCC4");
             act(() => {
                 result.current.addContextAttachment({
-                    name: "File.txt",
-                    icon,
+                    context: "file1",
+                    displayName: "Button.tsx",
+                    type: "file",
+                });
+                result.current.addContextAttachment({
+                    context: "page1",
+                    displayName: "Dashboard",
+                    type: "page",
                 });
             });
-            expect(result.current.attachments[0].icon).toBe(icon);
+            expect(result.current.attachments[0].type).toBe("file");
+            expect(result.current.attachments[1].type).toBe("page");
         });
     });
     describe("Removing Context Attachments", () => {
         /**
-         * Should remove a specific attachment by ID
+         * Should remove a specific context attachment by ID and update the merged helper
          */
-        it("should remove context attachment by ID", () => {
+        it("should remove context attachment by ID", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
+            let attachmentId = "";
             act(() => {
-                result.current.addContextAttachment({
-                    name: "Button.tsx",
+                const attachment = result.current.addContextAttachment({
+                    context: "selectedFile",
+                    displayName: "Button.tsx",
                 });
+                attachmentId = attachment.id;
             });
             expect(result.current.attachments).toHaveLength(1);
-            const attachmentId = result.current.attachments[0].id;
+            // Wait for initial helper registration
+            await waitFor(() => {
+                expect(mockAddContextHelper).toHaveBeenCalledWith(CONTEXT_ATTACHMENTS_HELPER_KEY, expect.any(Function));
+            });
             act(() => {
                 result.current.removeContextAttachment(attachmentId);
             });
             expect(result.current.attachments).toHaveLength(0);
+            // Wait for useEffect to update the helper (should return null when empty)
+            await waitFor(() => {
+                const lastCall = mockAddContextHelper.mock.calls
+                    .filter((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)
+                    .pop();
+                if (lastCall) {
+                    const helperFn = lastCall[1];
+                    expect(helperFn()).toBeNull();
+                }
+            });
         });
         /**
          * Should only remove the specified attachment when multiple exist
          */
-        it("should only remove specified attachment", () => {
+        it("should only remove specified attachment", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
+            let firstId = "";
             act(() => {
-                result.current.addContextAttachment({ name: "First.tsx" });
-                result.current.addContextAttachment({ name: "Second.tsx" });
+                const first = result.current.addContextAttachment({
+                    context: "file1",
+                    displayName: "First.tsx",
+                });
+                result.current.addContextAttachment({
+                    context: "file2",
+                    displayName: "Second.tsx",
+                });
+                firstId = first.id;
             });
             expect(result.current.attachments).toHaveLength(2);
-            const firstId = result.current.attachments[0].id;
+            // Wait for initial helper registration
+            await waitFor(() => {
+                expect(mockAddContextHelper).toHaveBeenCalledWith(CONTEXT_ATTACHMENTS_HELPER_KEY, expect.any(Function));
+            });
             act(() => {
                 result.current.removeContextAttachment(firstId);
             });
             expect(result.current.attachments).toHaveLength(1);
-            expect(result.current.attachments[0].name).toBe("Second.tsx");
+            expect(result.current.attachments[0].displayName).toBe("Second.tsx");
+            // Wait for useEffect to update the helper with remaining attachment
+            await waitFor(() => {
+                const lastCall = mockAddContextHelper.mock.calls
+                    .filter((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)
+                    .pop();
+                if (lastCall) {
+                    const helperFn = lastCall[1];
+                    const helperResult = helperFn();
+                    expect(helperResult).toHaveLength(1);
+                    expect(helperResult[0].displayName).toBe("Second.tsx");
+                }
+            });
         });
         /**
          * Should handle removing non-existent attachment gracefully
          */
-        it("should handle removing non-existent attachment gracefully", () => {
+        it("should handle removing non-existent attachment gracefully", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
+            // Wait for initial helper registration (useEffect runs on mount)
+            await waitFor(() => {
+                expect(mockAddContextHelper).toHaveBeenCalledWith(CONTEXT_ATTACHMENTS_HELPER_KEY, expect.any(Function));
+            });
             expect(() => {
                 act(() => {
                     result.current.removeContextAttachment("non-existent-id");
                 });
             }).not.toThrow();
+            // The helper is registered on mount, so removeContextHelper will be called
+            // when the component unmounts or when attachments change, but not for
+            // removing a non-existent attachment since attachments didn't change
+            // However, the cleanup function from the initial useEffect registration
+            // may be called. Let's just verify the attachment list is still empty.
+            expect(result.current.attachments).toHaveLength(0);
         });
     });
-    describe("Clearing All Attachments", () => {
+    describe("Clearing All Context Attachments", () => {
         /**
-         * Should clear all attachments at once
+         * Should clear all context attachments and update the merged helper to return null
          */
-        it("should clear all context attachments", () => {
+        it("should clear all context attachments", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
             act(() => {
-                result.current.addContextAttachment({ name: "First.tsx" });
-                result.current.addContextAttachment({ name: "Second.tsx" });
-                result.current.addContextAttachment({ name: "Third.tsx" });
+                result.current.addContextAttachment({
+                    context: "file1",
+                    displayName: "First.tsx",
+                });
+                result.current.addContextAttachment({
+                    context: "file2",
+                    displayName: "Second.tsx",
+                });
+                result.current.addContextAttachment({
+                    context: "file3",
+                    displayName: "Third.tsx",
+                });
             });
             expect(result.current.attachments).toHaveLength(3);
+            // Wait for initial helper registration
+            await waitFor(() => {
+                expect(mockAddContextHelper).toHaveBeenCalledWith(CONTEXT_ATTACHMENTS_HELPER_KEY, expect.any(Function));
+            });
             act(() => {
                 result.current.clearContextAttachments();
             });
             expect(result.current.attachments).toHaveLength(0);
+            // Wait for useEffect to update the helper (should return null when empty)
+            await waitFor(() => {
+                const lastCall = mockAddContextHelper.mock.calls
+                    .filter((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)
+                    .pop();
+                if (lastCall) {
+                    const helperFn = lastCall[1];
+                    expect(helperFn()).toBeNull();
+                }
+            });
         });
         /**
          * Should handle clearing when no attachments exist
          */
-        it("should handle clearing when no attachments exist", () => {
+        it("should handle clearing when no attachments exist", async () => {
             const { result } = renderHook(() => useTamboContextAttachment(), {
                 wrapper: createWrapper(),
             });
@@ -215,413 +381,13 @@ describe("TamboContextAttachmentProvider", () => {
                 });
             }).not.toThrow();
             expect(result.current.attachments).toHaveLength(0);
-        });
-    });
-    describe("Context Helpers Integration", () => {
-        /**
-         * Should automatically register context helpers when attachments are added
-         */
-        it("should register context helpers for attachments", async () => {
-            const wrapper = ({ children }) => (React.createElement(TamboContextHelpersProvider, null,
-                React.createElement(TamboContextAttachmentProvider, null, children)));
-            const { result } = renderHook(() => ({
-                attachment: useTamboContextAttachment(),
-                helpers: useTamboContextHelpers(),
-            }), { wrapper });
-            // Add attachment
-            act(() => {
-                result.current.attachment.addContextAttachment({
-                    name: "Button.tsx",
-                    metadata: { filePath: "/src/Button.tsx" },
-                });
-            });
-            // Wait for effect to run
-            await waitFor(async () => {
-                const contexts = await result.current.helpers.getAdditionalContext();
-                expect(contexts.length).toBeGreaterThan(0);
-            });
-            const contexts = await result.current.helpers.getAdditionalContext();
-            const attachmentContext = contexts.find((c) => c.name.includes(result.current.attachment.attachments[0].id));
-            expect(attachmentContext).toBeDefined();
-            expect(attachmentContext?.context).toHaveProperty("selectedComponent");
-        });
-        /**
-         * Should unregister context helpers when attachments are removed
-         */
-        it("should unregister context helpers when attachments are removed", async () => {
-            const wrapper = ({ children }) => (React.createElement(TamboContextHelpersProvider, null,
-                React.createElement(TamboContextAttachmentProvider, null, children)));
-            const { result } = renderHook(() => ({
-                attachment: useTamboContextAttachment(),
-                helpers: useTamboContextHelpers(),
-            }), { wrapper });
-            // Add attachment
-            act(() => {
-                result.current.attachment.addContextAttachment({
-                    name: "Button.tsx",
-                });
-            });
-            // Wait for context helper to be registered
-            await waitFor(async () => {
-                const contexts = await result.current.helpers.getAdditionalContext();
-                expect(contexts.length).toBeGreaterThan(0);
-            });
-            const initialContexts = await result.current.helpers.getAdditionalContext();
-            const initialCount = initialContexts.length;
-            const attachmentId = result.current.attachment.attachments[0].id;
-            // Remove attachment
-            act(() => {
-                result.current.attachment.removeContextAttachment(attachmentId);
-            });
-            // Wait for context helper to be unregistered
-            await waitFor(async () => {
-                const contexts = await result.current.helpers.getAdditionalContext();
-                expect(contexts.length).toBeLessThan(initialCount);
-            });
-        });
-        /**
-         * Should use default context data structure when no custom function provided
-         */
-        it("should use default context data structure", async () => {
-            const wrapper = ({ children }) => (React.createElement(TamboContextHelpersProvider, null,
-                React.createElement(TamboContextAttachmentProvider, null, children)));
-            const { result } = renderHook(() => ({
-                attachment: useTamboContextAttachment(),
-                helpers: useTamboContextHelpers(),
-            }), { wrapper });
-            act(() => {
-                result.current.attachment.addContextAttachment({
-                    name: "Button.tsx",
-                    metadata: { filePath: "/src/Button.tsx", type: "component" },
-                });
-            });
-            await waitFor(async () => {
-                const contexts = await result.current.helpers.getAdditionalContext();
-                expect(contexts.length).toBeGreaterThan(0);
-            });
-            const contexts = await result.current.helpers.getAdditionalContext();
-            const attachmentContext = contexts.find((c) => c.name.includes(result.current.attachment.attachments[0].id));
-            expect(attachmentContext?.context).toMatchObject({
-                selectedComponent: {
-                    name: "Button.tsx",
-                    instruction: expect.stringContaining("Tambo interactable component"),
-                    filePath: "/src/Button.tsx",
-                    type: "component",
-                },
-            });
-        });
-        /**
-         * Should use custom getContextHelperData function when provided
-         */
-        it("should use custom getContextHelperData function", async () => {
-            const customGetContextHelperData = jest.fn(async (context) => ({
-                selectedFile: {
-                    name: context.name,
-                    path: context.metadata?.filePath,
-                    customField: "custom value",
-                },
-            }));
-            const wrapper = ({ children }) => (React.createElement(TamboContextHelpersProvider, null,
-                React.createElement(TamboContextAttachmentProvider, { getContextHelperData: customGetContextHelperData }, children)));
-            const { result } = renderHook(() => ({
-                attachment: useTamboContextAttachment(),
-                helpers: useTamboContextHelpers(),
-            }), { wrapper });
-            act(() => {
-                result.current.attachment.addContextAttachment({
-                    name: "Button.tsx",
-                    metadata: { filePath: "/src/Button.tsx" },
-                });
-            });
-            // Wait for context helper to be registered and called
-            await waitFor(async () => {
-                const contexts = await result.current.helpers.getAdditionalContext();
-                expect(contexts.length).toBeGreaterThan(0);
-            }, { timeout: 2000 });
-            expect(customGetContextHelperData).toHaveBeenCalled();
-            const contexts = await result.current.helpers.getAdditionalContext();
-            const attachmentContext = contexts.find((c) => c.name.includes(result.current.attachment.attachments[0].id));
-            expect(attachmentContext?.context).toEqual({
-                selectedFile: {
-                    name: "Button.tsx",
-                    path: "/src/Button.tsx",
-                    customField: "custom value",
-                },
-            });
-        });
-        /**
-         * Should update context helpers when getContextHelperData function changes
-         * This tests the bug fix for stale closure issue
-         */
-        it("should update existing context helpers when getContextHelperData changes", async () => {
-            // Create a wrapper component with state to manage the function
-            const TestWrapper = ({ children }) => {
-                const [version, setVersion] = React.useState("v1");
-                const getContextHelperData = React.useCallback(async (context) => ({
-                    version,
-                    name: context.name,
-                }), [version]);
-                // Expose setVersion for the test
-                TestWrapper.setVersion = setVersion;
-                return (React.createElement(TamboContextHelpersProvider, null,
-                    React.createElement(TamboContextAttachmentProvider, { getContextHelperData: getContextHelperData }, children)));
-            };
-            const { result } = renderHook(() => ({
-                attachment: useTamboContextAttachment(),
-                helpers: useTamboContextHelpers(),
-            }), { wrapper: TestWrapper });
-            // Add attachment with first version
-            act(() => {
-                result.current.attachment.addContextAttachment({
-                    name: "Button.tsx",
-                });
-            });
-            // Wait for context helper to be registered
-            await waitFor(async () => {
-                const contexts = await result.current.helpers.getAdditionalContext();
-                expect(contexts.length).toBeGreaterThan(0);
-            }, { timeout: 2000 });
-            // Verify v1 context
-            let contexts = await result.current.helpers.getAdditionalContext();
-            let attachmentContext = contexts.find((c) => c.name.includes(result.current.attachment.attachments[0].id));
-            expect(attachmentContext?.context).toMatchObject({ version: "v1" });
-            // Change the version which will trigger a new getContextHelperData function
-            act(() => {
-                TestWrapper.setVersion("v2");
-            });
-            // Wait for context to update to v2
-            await waitFor(async () => {
-                const contexts = await result.current.helpers.getAdditionalContext();
-                const context = contexts.find((c) => c.name.includes(result.current.attachment.attachments[0].id));
-                return context?.context.version === "v2";
-            }, { timeout: 2000 });
-            // Verify v2 context - should be updated, not stale
-            contexts = await result.current.helpers.getAdditionalContext();
-            attachmentContext = contexts.find((c) => c.name.includes(result.current.attachment.attachments[0].id));
-            expect(attachmentContext?.context).toMatchObject({ version: "v2" });
-        });
-        /**
-         * Should handle errors in getContextHelperData gracefully
-         */
-        it("should handle errors in getContextHelperData gracefully", async () => {
-            const consoleErrorSpy = jest
-                .spyOn(console, "error")
-                .mockImplementation(() => { });
-            const errorGetData = jest.fn(async () => {
-                throw new Error("Custom data error");
-            });
-            const wrapper = ({ children }) => (React.createElement(TamboContextHelpersProvider, null,
-                React.createElement(TamboContextAttachmentProvider, { getContextHelperData: errorGetData }, children)));
-            const { result } = renderHook(() => ({
-                attachment: useTamboContextAttachment(),
-                helpers: useTamboContextHelpers(),
-            }), { wrapper });
-            act(() => {
-                result.current.attachment.addContextAttachment({
-                    name: "Button.tsx",
-                });
-            });
-            // Wait for effect to run and try to call the function
+            // Wait for useEffect to run - helper should be registered and return null
             await waitFor(() => {
-                expect(result.current.attachment.attachments.length).toBeGreaterThan(0);
-            }, { timeout: 2000 });
-            // Try to get contexts, which will trigger the error
-            await result.current.helpers.getAdditionalContext();
-            expect(errorGetData).toHaveBeenCalled();
-            expect(consoleErrorSpy).toHaveBeenCalled();
-            consoleErrorSpy.mockRestore();
-        });
-    });
-    describe("Custom Suggestions", () => {
-        /**
-         * Should start with null custom suggestions
-         */
-        it("should initialize with null custom suggestions", () => {
-            const { result } = renderHook(() => useTamboContextAttachment(), {
-                wrapper: createWrapper(),
-            });
-            expect(result.current.customSuggestions).toBeNull();
-        });
-        /**
-         * Should set custom suggestions
-         */
-        it("should set custom suggestions", () => {
-            const { result } = renderHook(() => useTamboContextAttachment(), {
-                wrapper: createWrapper(),
-            });
-            const suggestions = [
-                {
-                    id: "1",
-                    title: "Edit component",
-                    detailedSuggestion: "Modify the Button component",
-                    messageId: "",
-                },
-                {
-                    id: "2",
-                    title: "Add feature",
-                    detailedSuggestion: "Add a new feature",
-                    messageId: "",
-                },
-            ];
-            act(() => {
-                result.current.setCustomSuggestions(suggestions);
-            });
-            expect(result.current.customSuggestions).toEqual(suggestions);
-        });
-        /**
-         * Should clear custom suggestions by setting to null
-         */
-        it("should clear custom suggestions", () => {
-            const { result } = renderHook(() => useTamboContextAttachment(), {
-                wrapper: createWrapper(),
-            });
-            const suggestions = [
-                {
-                    id: "1",
-                    title: "Test",
-                    detailedSuggestion: "Test suggestion",
-                    messageId: "",
-                },
-            ];
-            act(() => {
-                result.current.setCustomSuggestions(suggestions);
-            });
-            expect(result.current.customSuggestions).toEqual(suggestions);
-            act(() => {
-                result.current.setCustomSuggestions(null);
-            });
-            expect(result.current.customSuggestions).toBeNull();
-        });
-        /**
-         * Should update custom suggestions when changed
-         */
-        it("should update custom suggestions", () => {
-            const { result } = renderHook(() => useTamboContextAttachment(), {
-                wrapper: createWrapper(),
-            });
-            const firstSuggestions = [
-                {
-                    id: "1",
-                    title: "First",
-                    detailedSuggestion: "First suggestion",
-                    messageId: "",
-                },
-            ];
-            const secondSuggestions = [
-                {
-                    id: "2",
-                    title: "Second",
-                    detailedSuggestion: "Second suggestion",
-                    messageId: "",
-                },
-            ];
-            act(() => {
-                result.current.setCustomSuggestions(firstSuggestions);
-            });
-            expect(result.current.customSuggestions).toEqual(firstSuggestions);
-            act(() => {
-                result.current.setCustomSuggestions(secondSuggestions);
-            });
-            expect(result.current.customSuggestions).toEqual(secondSuggestions);
-        });
-    });
-    describe("Combined Workflows", () => {
-        /**
-         * Should handle adding attachment and setting custom suggestions together
-         */
-        it("should handle attachment and custom suggestions together", () => {
-            const { result } = renderHook(() => useTamboContextAttachment(), {
-                wrapper: createWrapper(),
-            });
-            const suggestions = [
-                {
-                    id: "1",
-                    title: "Edit file",
-                    detailedSuggestion: "Edit this file",
-                    messageId: "",
-                },
-            ];
-            act(() => {
-                result.current.addContextAttachment({
-                    name: "Button.tsx",
-                    metadata: { filePath: "/src/Button.tsx" },
-                });
-                result.current.setCustomSuggestions(suggestions);
-            });
-            expect(result.current.attachments).toHaveLength(1);
-            expect(result.current.customSuggestions).toEqual(suggestions);
-        });
-        /**
-         * Should clear suggestions when clearing attachments
-         */
-        it("should independently manage attachments and suggestions", () => {
-            const { result } = renderHook(() => useTamboContextAttachment(), {
-                wrapper: createWrapper(),
-            });
-            const suggestions = [
-                {
-                    id: "1",
-                    title: "Test",
-                    detailedSuggestion: "Test suggestion",
-                    messageId: "",
-                },
-            ];
-            act(() => {
-                result.current.addContextAttachment({ name: "File.tsx" });
-                result.current.setCustomSuggestions(suggestions);
-            });
-            // Clear attachments
-            act(() => {
-                result.current.clearContextAttachments();
-            });
-            // Suggestions should remain
-            expect(result.current.attachments).toHaveLength(0);
-            expect(result.current.customSuggestions).toEqual(suggestions);
-            // Can clear suggestions separately
-            act(() => {
-                result.current.setCustomSuggestions(null);
-            });
-            expect(result.current.customSuggestions).toBeNull();
-        });
-    });
-    describe("Cleanup", () => {
-        /**
-         * Should cleanup context helpers on unmount
-         */
-        it("should cleanup context helpers on unmount", async () => {
-            // Use a shared provider wrapper so context helpers persist
-            const SharedWrapper = ({ children }) => (React.createElement(TamboContextHelpersProvider, null, children));
-            // First render with attachment provider
-            const { result: attachmentResult, unmount } = renderHook(() => ({
-                attachment: useTamboContextAttachment(),
-                helpers: useTamboContextHelpers(),
-            }), {
-                wrapper: ({ children }) => (React.createElement(SharedWrapper, null,
-                    React.createElement(TamboContextAttachmentProvider, null, children))),
-            });
-            // Add attachment
-            act(() => {
-                attachmentResult.current.attachment.addContextAttachment({
-                    name: "Button.tsx",
-                });
+                expect(mockAddContextHelper).toHaveBeenCalledWith(CONTEXT_ATTACHMENTS_HELPER_KEY, expect.any(Function));
             });
-            const attachmentId = attachmentResult.current.attachment.attachments[0].id;
-            // Wait for context helper to be registered
-            await waitFor(async () => {
-                const contexts = await attachmentResult.current.helpers.getAdditionalContext();
-                const hasContext = contexts.some((c) => c.name.includes(attachmentId));
-                expect(hasContext).toBe(true);
-            });
-            // Unmount the attachment provider
-            unmount();
-            // Create new hook to check cleanup
-            const { result: helpersResult } = renderHook(() => useTamboContextHelpers(), { wrapper: SharedWrapper });
-            // Wait a bit for cleanup effect to run
-            await new Promise((resolve) => setTimeout(resolve, 100));
-            const contexts = await helpersResult.current.getAdditionalContext();
-            // Should not contain the attachment context after unmount
-            const hasAttachmentContext = contexts.some((c) => c.name.includes(attachmentId));
-            expect(hasAttachmentContext).toBe(false);
+            const helperFn = mockAddContextHelper.mock.calls.find((call) => call[0] === CONTEXT_ATTACHMENTS_HELPER_KEY)?.[1];
+            expect(helperFn).toBeDefined();
+            expect(helperFn()).toBeNull();
         });
     });
 });