Atomic assistant #951

browser/data-browser/package.json

@@ -14,19 +14,23 @@
     "@dnd-kit/utilities": "^3.2.2",
     "@emoji-mart/react": "^1.1.1",
     "@emotion/is-prop-valid": "^1.3.1",
+    "@modelcontextprotocol/sdk": "^1.11.4",
+    "@openrouter/ai-sdk-provider": "^0.4.3",
     "@radix-ui/react-popover": "^1.1.2",
     "@radix-ui/react-scroll-area": "^1.2.0",
     "@radix-ui/react-tabs": "^1.1.1",
     "@tanstack/react-router": "^1.95.1",
-    "@tiptap/extension-image": "^2.9.1",
-    "@tiptap/extension-link": "^2.9.1",
-    "@tiptap/extension-placeholder": "^2.9.1",
-    "@tiptap/extension-typography": "^2.9.1",
-    "@tiptap/pm": "^2.9.1",
-    "@tiptap/react": "^2.9.1",
-    "@tiptap/starter-kit": "^2.9.1",
-    "@tiptap/suggestion": "^2.9.1",
+    "@tiptap/extension-image": "^2.11.7",
+    "@tiptap/extension-link": "^2.11.7",
+    "@tiptap/extension-mention": "^2.11.7",
+    "@tiptap/extension-placeholder": "^2.11.7",
+    "@tiptap/extension-typography": "^2.11.7",
+    "@tiptap/pm": "^2.11.7",
+    "@tiptap/react": "^2.11.7",
+    "@tiptap/starter-kit": "^2.11.7",
+    "@tiptap/suggestion": "^2.11.7",
     "@tomic/react": "workspace:*",
+    "ai": "^4.3.16",
     "emoji-mart": "^5.6.0",
     "polished": "^4.3.1",
     "prismjs": "^1.29.0",
@@ -50,7 +54,8 @@
     "styled-components": "^6.1.13",
     "stylis": "4.3.0",
     "tippy.js": "^6.3.7",
-    "tiptap-markdown": "^0.8.10"
+    "tiptap-markdown": "^0.8.10",
+    "zod": "^3.24.2"
   },
   "devDependencies": {
     "@tanstack/router-devtools": "^1.95.1",

browser/data-browser/src/Providers.tsx

@@ -18,6 +18,7 @@ import CrashPage from './views/CrashPage';
 import { AppSettingsContextProvider } from './helpers/AppSettings';
 import { NavStateProvider } from './components/NavState';
 import { Toaster } from './components/Toaster';
+import { McpServersProvider } from './components/AI/MCP/useMcpServers';
 
 // Setup bugsnag for error handling, but only if there's an API key
 const ErrBoundary = window.bugsnagApiKey
@@ -39,37 +40,39 @@ export const Providers: React.FC<React.PropsWithChildren> = ({ children }) => {
   return (
     <NavStateProvider>
       <AppSettingsContextProvider>
-        <ControlLockProvider>
-          <HotKeysWrapper>
-            <StyleSheetManager shouldForwardProp={shouldForwardProp}>
-              <ThemeWrapper>
-                <GlobalStyle />
-                <ErrBoundary FallbackComponent={CrashPage}>
-                  {/* Default form validation provider. Does not do anything on its own but will make sure useValidation works without context*/}
-                  <FormValidationContextProvider
-                    onValidationChange={() => undefined}
-                  >
-                    <Toaster />
-                    <MetaSetter />
-                    <DropdownContainer>
-                      <DialogGlobalContextProvider>
-                        <PopoverContainer>
-                          <DropdownContainer>
-                            <NewResourceUIProvider>
-                              <SkipNav />
-                              <NavWrapper>{children}</NavWrapper>
-                            </NewResourceUIProvider>
-                          </DropdownContainer>
-                        </PopoverContainer>
-                        <NetworkIndicator />
-                      </DialogGlobalContextProvider>
-                    </DropdownContainer>
-                  </FormValidationContextProvider>
-                </ErrBoundary>
-              </ThemeWrapper>
-            </StyleSheetManager>
-          </HotKeysWrapper>
-        </ControlLockProvider>
+        <McpServersProvider>
+          <ControlLockProvider>
+            <HotKeysWrapper>
+              <StyleSheetManager shouldForwardProp={shouldForwardProp}>
+                <ThemeWrapper>
+                  <GlobalStyle />
+                  <ErrBoundary FallbackComponent={CrashPage}>
+                    {/* Default form validation provider. Does not do anything on its own but will make sure useValidation works without context*/}
+                    <FormValidationContextProvider
+                      onValidationChange={() => undefined}
+                    >
+                      <Toaster />
+                      <MetaSetter />
+                      <DropdownContainer>
+                        <DialogGlobalContextProvider>
+                          <PopoverContainer>
+                            <DropdownContainer>
+                              <NewResourceUIProvider>
+                                <SkipNav />
+                                <NavWrapper>{children}</NavWrapper>
+                              </NewResourceUIProvider>
+                            </DropdownContainer>
+                          </PopoverContainer>
+                          <NetworkIndicator />
+                        </DialogGlobalContextProvider>
+                      </DropdownContainer>
+                    </FormValidationContextProvider>
+                  </ErrBoundary>
+                </ThemeWrapper>
+              </StyleSheetManager>
+            </HotKeysWrapper>
+          </ControlLockProvider>
+        </McpServersProvider>
       </AppSettingsContextProvider>
     </NavStateProvider>
   );

browser/data-browser/src/chunks/MarkdownEditor/AIChatInput/AsyncAIChatInput.tsx

@@ -0,0 +1,222 @@
+import { EditorContent, useEditor, type JSONContent } from '@tiptap/react';
+import { styled } from 'styled-components';
+import StarterKit from '@tiptap/starter-kit';
+import Mention from '@tiptap/extension-mention';
+import { TiptapContextProvider } from '../TiptapContext';
+import { EditorWrapperBase } from '../EditorWrapperBase';
+import { searchSuggestionBuilder } from './resourceSuggestions';
+import { useEffect, useRef, useState } from 'react';
+import { EditorEvents } from '../EditorEvents';
+import { Markdown } from 'tiptap-markdown';
+import { useStore } from '@tomic/react';
+import { useSettings } from '../../../helpers/AppSettings';
+import type { Node } from '@tiptap/pm/model';
+import Placeholder from '@tiptap/extension-placeholder';
+import { useMcpServers } from '../../../components/AI/MCP/useMcpServers';
+import type {
+  AtomicResourceSuggestion,
+  MCPResourceSuggestion,
+  MentionItem,
+} from './types';
+import { Row } from '../../../components/Row';
+import {
+  IconButton,
+  IconButtonVariant,
+} from '../../../components/IconButton/IconButton';
+import { FaArrowRight } from 'react-icons/fa6';
+
+const createAttribute = (propName: string, dataName: string) => {
+  return {
+    [propName]: {
+      default: null,
+      parseHTML: (element: HTMLElement) => element.getAttribute(dataName),
+      renderHTML: (attributes: Record<string, unknown>) => {
+        if (!attributes[propName]) {
+          return {};
+        }
+
+        return {
+          [dataName]: attributes[propName],
+        };
+      },
+    },
+  };
+};
+
+// Modify the Mention extension to allow serializing to markdown.
+const SerializableMention = Mention.extend({
+  addStorage() {
+    return {
+      markdown: {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        serialize(state: any, node: Node) {
+          state.write('@' + (node.attrs.label || ''));
+          state.renderContent(node);
+          state.flushClose(1);
+        },
+      },
+    };
+  },
+  addAttributes() {
+    return {
+      ...createAttribute('type', 'data-type'),
+      ...createAttribute('serverId', 'data-server-id'),
+      ...createAttribute('mimeType', 'data-mime-type'),
+      ...createAttribute('id', 'data-id'),
+      ...createAttribute('label', 'data-label'),
+      ...createAttribute('isA', 'data-is-a'),
+    };
+  },
+});
+
+interface AsyncAIChatInputProps {
+  onMentionUpdate: (mentions: MentionItem[]) => void;
+  onChange: (markdown: string) => void;
+  onSubmit: () => void;
+  hasFiles: boolean;
+}
+
+const AsyncAIChatInput: React.FC<
+  React.PropsWithChildren<AsyncAIChatInputProps>
+> = ({ onMentionUpdate, onChange, onSubmit, children, hasFiles }) => {
+  const store = useStore();
+  const { drive, mcpServers } = useSettings();
+  const [markdown, setMarkdown] = useState('');
+  const markdownRef = useRef(markdown);
+  const onSubmitRef = useRef(onSubmit);
+  const { serversWithResources, searchResourcesOfServer } = useMcpServers();
+
+  const editor = useEditor(
+    {
+      extensions: [
+        Markdown.configure({
+          html: true,
+        }),
+        StarterKit.extend({
+          addKeyboardShortcuts() {
+            return {
+              Enter: () => {
+                // Check if the cursor is in a code block, if so allow the user to press enter.
+                // Pressing shift + enter will exit the code block.
+                if ('language' in this.editor.getAttributes('codeBlock')) {
+                  return false;
+                }
+
+                // The content has to be read from a ref because this callback is not updated often leading to stale content.
+                onSubmitRef.current();
+                setMarkdown('');
+                this.editor.commands.clearContent();
+
+                return true;
+              },
+            };
+          },
+        }).configure({
+          blockquote: false,
+          bulletList: false,
+          orderedList: false,
+          heading: false,
+          listItem: false,
+          horizontalRule: false,
+          bold: false,
+          strike: false,
+          italic: false,
+        }),
+        SerializableMention.configure({
+          HTMLAttributes: {
+            class: 'ai-chat-mention',
+          },
+          suggestion: searchSuggestionBuilder(
+            store,
+            drive,
+            mcpServers.filter(server =>
+              serversWithResources.includes(server.id),
+            ),
+            searchResourcesOfServer,
+          ),
+          renderText({ options, node }) {
+            return `${options.suggestion.char}bla${node.attrs.title}`;
+          },
+        }),
+        Placeholder.configure({
+          placeholder: 'Ask me anything...',
+        }),
+      ],
+      autofocus: true,
+    },
+    [serversWithResources, searchResourcesOfServer],
+  );
+
+  const handleChange = (value: string) => {
+    setMarkdown(value);
+    markdownRef.current = value;
+    onChange(value);
+
+    if (!editor) {
+      return;
+    }
+
+    const mentions = digForMentions(editor.getJSON());
+    onMentionUpdate(mentions);
+  };
+
+  useEffect(() => {
+    markdownRef.current = markdown;
+    onSubmitRef.current = onSubmit;
+  }, [markdown, onSubmit]);
+
+  return (
+    <>
+      <EditorWrapper hideEditor={false}>
+        <TiptapContextProvider editor={editor}>
+          <EditorContent editor={editor} />
+          <EditorEvents onChange={handleChange} />
+        </TiptapContextProvider>
+      </EditorWrapper>
+      <Row justify='space-between'>
+        {children}
+        <IconButton
+          disabled={markdown.length === 0 && !hasFiles}
+          onClick={() => {
+            onSubmit();
+            setMarkdown('');
+            editor?.commands.clearContent();
+          }}
+          title='Send'
+          variant={IconButtonVariant.Fill}
+        >
+          <FaArrowRight />
+        </IconButton>
+      </Row>
+    </>
+  );
+};
+
+export default AsyncAIChatInput;
+
+const EditorWrapper = styled(EditorWrapperBase)`
+  padding: ${p => p.theme.size(2)};
+  font-size: 16px;
+  line-height: 1.5;
+
+  .ai-chat-mention {
+    background-color: ${p => p.theme.colors.mainSelectedBg};
+    color: ${p => p.theme.colors.mainSelectedFg};
+    border-radius: 5px;
+    padding-inline: ${p => p.theme.size(1)};
+  }
+`;
+
+function digForMentions(
+  data: JSONContent,
+): Array<MCPResourceSuggestion | AtomicResourceSuggestion> {
+  if (data.type === 'mention') {
+    return [data.attrs as MCPResourceSuggestion | AtomicResourceSuggestion];
+  }
+
+  if (data.content) {
+    return data.content.flatMap(digForMentions);
+  }
+
+  return [];
+}