AT-360 Added Diagram Chat Functionality To Mermaid/sdk So MCP Server Can Directly Access Mermaid AI

packages/sdk/src/index.e2e.test.ts

@@ -180,7 +180,7 @@ describe('deleteDocument', () => {
     expect(deletedDoc.projectID).toStrictEqual(newDocument.projectID);
 
     expect(await client.getDocuments(testProjectId)).not.toContainEqual(newDocument);
-  });
+  }, 20000); // 20 seconds — test makes 4 sequential HTTP calls
 });
 
 describe('getDocument', () => {

packages/sdk/src/index.test.ts

@@ -87,6 +87,47 @@ describe('MermaidChart', () => {
     });
   });
 
+  describe('#diagramChat', () => {
+    beforeEach(async () => {
+      await client.setAccessToken('test-access-token');
+    });
+
+    it('should parse stream body into text and documentChatThreadID', async () => {
+      const streamBody = [
+        '0:"Hello, "',
+        '0:"here is your diagram!"',
+        '2:[{"documentChatThreadID":"thread-abc-123"}]',
+        'e:{"finishReason":"stop"}',
+        'd:{"finishReason":"stop"}',
+      ].join('\n');
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      vi.spyOn((client as any).axios, 'post').mockResolvedValue({ data: streamBody });
+
+      const result = await client.diagramChat({
+        message: 'Create a flowchart',
+        documentID: 'doc-123',
+      });
+
+      expect(result.text).toBe('Hello, here is your diagram!');
+      expect(result.documentChatThreadID).toBe('thread-abc-123');
+      expect(result.documentID).toBe('doc-123');
+    });
+
+    it('should throw AICreditsLimitExceededError on 402', async () => {
+      vi.spyOn(client, 'diagramChat').mockRejectedValue(
+        new AICreditsLimitExceededError('AI credits limit exceeded'),
+      );
+
+      await expect(
+        client.diagramChat({
+          message: 'Create a flowchart',
+          documentID: 'doc-123',
+        }),
+      ).rejects.toThrow(AICreditsLimitExceededError);
+    });
+  });
+
   describe('#repairDiagram', () => {
     beforeEach(async () => {
       await client.setAccessToken('test-access-token');

packages/sdk/src/index.ts

@@ -10,6 +10,8 @@ import {
 import type {
   AuthState,
   AuthorizationData,
+  DiagramChatRequest,
+  DiagramChatResponse,
   Document,
   InitParams,
   MCDocument,
@@ -24,6 +26,63 @@ import { URLS } from './urls.js';
 const defaultBaseURL = 'https://www.mermaid.ai'; // "http://127.0.0.1:5174"
 const authorizationURLTimeout = 60_000;
 
+/**
+ * Parses text tokens from a Vercel AI SDK data-stream response body.
+ *
+ * The stream format uses line prefixes:
+ *   `0:"text_chunk"\n` – text token (JSON-encoded string)
+ *   `2:[{"documentChatThreadID":"thread-abc-123"}]\n`      – data payload (JSON-encoded array)
+ *   `e:{...}\n`        – step finish
+ *   `d:{...}\n`        – stream done
+ */
+function parseVercelAIStreamText(rawBody: string): string {
+  return rawBody
+    .split('\n')
+    .filter((line) => line.startsWith('0:'))
+    .map((line) => {
+      try {
+        const value = JSON.parse(line.slice(2));
+        return typeof value === 'string' ? value : '';
+      } catch {
+        return '';
+      }
+    })
+    .join('');
+}
+
+/**
+ * Extracts data payloads from a Vercel AI SDK data-stream response body.
+ * Returns the first `documentChatThreadID` found in the stream, if any.
+ */
+function parseVercelAIStreamData(rawBody: string): { documentChatThreadID?: string } {
+  let documentChatThreadID: string | undefined;
+
+  for (const line of rawBody.split('\n')) {
+    if (!line.startsWith('2:')) {
+      continue;
+    }
+    try {
+      const items: unknown[] = JSON.parse(line.slice(2));
+      for (const item of items) {
+        if (item && typeof item === 'object' && 'documentChatThreadID' in item) {
+          const value = (item as Record<string, unknown>).documentChatThreadID;
+          if (typeof value === 'string') {
+            documentChatThreadID = value;
+            break;
+          }
+        }
+      }
+    } catch {
+      // ignore malformed lines
+    }
+    if (documentChatThreadID) {
+      break;
+    }
+  }
+
+  return { documentChatThreadID };
+}
+
 export class MermaidChart {
   private clientID: string;
   #baseURL!: string;
@@ -319,4 +378,79 @@ export class MermaidChart {
       throw error;
     }
   }
+
+  /**
+   * Chat with Mermaid AI about a diagram.
+   *
+   * Sends a single user message to the Mermaid AI chat endpoint.  The backend
+   * automatically fetches the full conversation history from the database
+   * (when `documentChatThreadID` is provided), so callers never need to track
+   * or resend previous messages.
+   *
+   * @param request - The chat request containing the user message and diagram context
+   * @returns The AI response text and the chat thread ID
+   * @throws {@link AICreditsLimitExceededError} if AI credits limit is exceeded (HTTP 402)
+   */
+  public async diagramChat(request: DiagramChatRequest): Promise<DiagramChatResponse> {
+    const { message, documentID , code = '', documentChatThreadID } = request;
+
+    // Send only the current user message. The backend will prepend the stored
+    // conversation history when autoFetchHistory is true (see AIChatRequestData).
+    const messages = [
+      {
+        id: uuid(),
+        role: 'user' as const,
+        content: message,
+        experimental_attachments: [] as [],
+      },
+    ];
+
+    const requestBody = {
+      messages,
+      code,
+      documentID,
+      documentChatThreadID,
+      // parentID null: the backend already handles finding the correct parent
+      parentID: null,
+      // Tell the backend to fetch DB history and prepend it before calling the AI.
+      autoFetchHistory: true,
+    };
+
+    try {
+      // responseType: 'text' buffers the full stream body as a plain string so we
+      // can parse the Vercel AI SDK data-stream format after the request completes.
+      const response = await this.axios.post<string>(URLS.rest.openai.chat, requestBody, {
+        responseType: 'text',
+        timeout: 120_000,
+      });
+
+      const rawBody = response.data;
+      const text = parseVercelAIStreamText(rawBody);
+      const { documentChatThreadID: returnedThreadID } = parseVercelAIStreamData(rawBody);
+
+      return {
+        text,
+        documentChatThreadID: returnedThreadID ?? documentChatThreadID,
+        documentID,
+      };
+    } catch (error: unknown) {
+      if (
+        error &&
+        typeof error === 'object' &&
+        'response' in error &&
+        error.response &&
+        typeof error.response === 'object' &&
+        'status' in error.response &&
+        (error as { response: { status: number } }).response.status === 402
+      ) {
+        const axiosError = error as { response: { status: number; data?: unknown } };
+        throw new AICreditsLimitExceededError(
+          typeof axiosError.response.data === 'string'
+            ? axiosError.response.data
+            : 'AI credits limit exceeded',
+        );
+      }
+      throw error;
+    }
+  }
 }

packages/sdk/src/types.ts

@@ -96,6 +96,49 @@ export interface RepairDiagramRequest {
   userID?: string;
 }
 
+/**
+ * Request parameters for chatting with the Mermaid AI about a diagram.
+ */
+export interface DiagramChatRequest {
+  /** The user's chat message / question */
+  message: string;
+  /**
+   * MermaidChart document ID to associate the chat thread with.
+   */
+  documentID: string;
+  /**
+   * Optional Mermaid diagram code to use as context.
+   * Pass the current diagram code when the user wants to modify or discuss a specific
+   * diagram. Leave empty (or omit) for general Mermaid questions.
+   * Defaults to an empty string.
+   */
+  code?: string;
+  /**
+   * Existing chat thread ID to continue a conversation.
+   * Returned from a previous diagramChat() call.
+   * When provided, the backend automatically fetches the stored conversation
+   * history from the database so the AI has full context.
+   */
+  documentChatThreadID?: string;
+}
+
+/**
+ * Response from chatting with the Mermaid AI.
+ */
+export interface DiagramChatResponse {
+  /** The full AI-generated response text (may contain Mermaid code blocks) */
+  text: string;
+  /**
+   * The chat thread ID created or used for this conversation.
+   * Pass this back as documentChatThreadID in subsequent calls to continue the thread.
+   */
+  documentChatThreadID?: string;
+  /**
+   * The document ID used for this conversation. Same as the one passed in the request.
+   */
+  documentID?: string;
+}
+
 /**
  * Response from repairing a diagram.
  * Matches OpenAIGenerationResult from collab.

packages/sdk/src/urls.ts

@@ -41,6 +41,7 @@ export const URLS = {
     },
     openai: {
       repair: `/rest-api/openai/repair`,
+      chat: `/rest-api/openai/chat`,
     },
   },
   raw: (document: Pick<MCDocument, 'documentID' | 'major' | 'minor'>, theme: 'light' | 'dark') => {