fix(docs): stabilize ai chat tool handling

Author: junghyeonsu

docs/app/api/chat/route.ts

@@ -1,5 +1,5 @@
 import { createOpenAI } from "@ai-sdk/openai";
-import { stepCountIs, streamText } from "ai";
+import { convertToModelMessages, safeValidateUIMessages, stepCountIs, streamText } from "ai";
 import { systemPrompt } from "@/lib/ai/system-prompt";
 import { clientTools } from "@/lib/ai/tools";
 import { getMCPTools } from "@/lib/ai/mcp-client";
@@ -15,16 +15,10 @@ const llmRouter = createOpenAI({
   name: "llm-router",
 });
 
+const llmRouterModel = process.env.LLM_ROUTER_MODEL?.trim() || "openai/gpt-4o";
+
 const chatRequestSchema = z.object({
-  messages: z
-    .array(
-      z
-        .object({
-          role: z.string().min(1),
-        })
-        .passthrough(),
-    )
-    .min(1),
+  messages: z.array(z.unknown()).min(1),
 });
 
 export async function POST(req: Request) {
@@ -40,18 +34,30 @@ export async function POST(req: Request) {
     return Response.json({ error: "Invalid request body" }, { status: 400 });
   }
 
-  const messages = parsed.data.messages as Parameters<typeof streamText>[0]["messages"];
-
   const mcpTools = await getMCPTools();
+  const tools = {
+    ...clientTools,
+    ...mcpTools,
+  };
+
+  const validatedMessages = await safeValidateUIMessages({
+    messages: parsed.data.messages,
+  });
+
+  if (!validatedMessages.success) {
+    return Response.json({ error: "Invalid messages format" }, { status: 400 });
+  }
+
+  const messages = await convertToModelMessages(validatedMessages.data, {
+    tools,
+    ignoreIncompleteToolCalls: true,
+  });
 
   const result = streamText({
-    model: llmRouter(process.env.LLM_ROUTER_MODEL ?? "openai/gpt-4o"),
+    model: llmRouter(llmRouterModel),
     system: systemPrompt,
     messages,
-    tools: {
-      ...clientTools,
-      ...mcpTools,
-    },
+    tools,
     stopWhen: stepCountIs(5),
   });
 

docs/lib/ai/mcp-client.ts

@@ -20,41 +20,108 @@ interface MCPToolDefinition {
 
 interface JSONRPCResponse {
   jsonrpc: string;
-  id: number;
+  id: number | string | null;
   result?: {
     tools?: MCPToolDefinition[];
     content?: Array<{ type: string; text?: string }>;
   };
   error?: { code: number; message: string };
 }
 
+interface JSONRPCRequest {
+  jsonrpc: "2.0";
+  method: string;
+  params?: Record<string, unknown>;
+  id?: number;
+}
+
+const MCP_PROTOCOL_VERSION = "2025-03-26";
+
 let cachedSessionId: string | null = null;
+let isInitialized = false;
+let initializingPromise: Promise<void> | null = null;
+
+function parseSSEJSONResponse(rawBody: string, method: string): JSONRPCResponse {
+  const events = rawBody.split(/\r?\n\r?\n/);
+
+  for (const eventBlock of events) {
+    const dataLines = eventBlock
+      .split(/\r?\n/)
+      .filter((line) => line.startsWith("data:"))
+      .map((line) => line.slice(5).trimStart())
+      .filter(Boolean);
+
+    if (dataLines.length === 0) {
+      continue;
+    }
+
+    const payload = dataLines.join("\n").trim();
+    if (!payload || payload === "[DONE]") {
+      continue;
+    }
+
+    try {
+      return JSON.parse(payload) as JSONRPCResponse;
+    } catch {
+      // ignore non-JSON SSE frames and continue
+    }
+  }
+
+  throw new Error(`Invalid MCP SSE response for "${method}"`);
+}
+
+function parseMCPResponseBody(
+  rawBody: string,
+  method: string,
+  contentType: string | null,
+): JSONRPCResponse {
+  if (!rawBody.trim()) {
+    return { jsonrpc: "2.0", id: null };
+  }
+
+  if (contentType?.includes("text/event-stream") || rawBody.trimStart().startsWith("event:")) {
+    return parseSSEJSONResponse(rawBody, method);
+  }
+
+  try {
+    return JSON.parse(rawBody) as JSONRPCResponse;
+  } catch (error) {
+    throw new Error(`Invalid MCP JSON response for "${method}": ${(error as Error).message}`);
+  }
+}
 
 async function mcpRequest(
   method: string,
   params: Record<string, unknown> = {},
+  options: { notification?: boolean } = {},
 ): Promise<JSONRPCResponse> {
   const url = process.env.SEED_DOCS_MCP_SERVER_URL;
   if (!url) throw new Error("SEED_DOCS_MCP_SERVER_URL is not set");
 
   const headers: Record<string, string> = {
     "Content-Type": "application/json",
-    Accept: "application/json",
+    Accept: "application/json, text/event-stream",
+    "MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
   };
 
   if (cachedSessionId) {
     headers["Mcp-Session-Id"] = cachedSessionId;
   }
 
+  const body: JSONRPCRequest = {
+    jsonrpc: "2.0",
+    method,
+    params,
+  };
+
+  if (!options.notification) {
+    body.id = Date.now();
+  }
+
   const response = await fetch(url, {
     method: "POST",
     headers,
-    body: JSON.stringify({
-      jsonrpc: "2.0",
-      id: Date.now(),
-      method,
-      params,
-    }),
+    body: JSON.stringify(body),
   });
 
   // 세션 ID 저장
@@ -63,24 +130,55 @@ async function mcpRequest(
     cachedSessionId = sessionId;
   }
 
+  const rawBody = await response.text();
+  const contentType = response.headers.get("content-type");
+
   if (!response.ok) {
-    throw new Error(`MCP request failed: ${response.status} ${response.statusText}`);
+    const details = rawBody.trim();
+    const suffix = details ? ` - ${details.slice(0, 300)}` : "";
+    throw new Error(`MCP request failed: ${response.status} ${response.statusText}${suffix}`);
   }
 
-  return response.json();
+  // notification 응답 혹은 빈 본문(202/204) 처리
+  if (options.notification || !rawBody.trim()) {
+    return { jsonrpc: "2.0", id: null };
+  }
+
+  return parseMCPResponseBody(rawBody, method, contentType);
+}
+
+async function mcpNotification(
+  method: string,
+  params: Record<string, unknown> = {},
+): Promise<void> {
+  await mcpRequest(method, params, { notification: true });
 }
 
 /**
  * MCP 서버 초기화 (initialize + initialized 핸드셰이크)
  */
 async function initializeMCP(): Promise<void> {
-  await mcpRequest("initialize", {
-    protocolVersion: "2025-03-26",
-    capabilities: {},
-    clientInfo: { name: "seed-docs-ai", version: "1.0.0" },
+  if (isInitialized) return;
+  if (initializingPromise) return initializingPromise;
+
+  initializingPromise = (async () => {
+    const initResponse = await mcpRequest("initialize", {
+      protocolVersion: MCP_PROTOCOL_VERSION,
+      capabilities: {},
+      clientInfo: { name: "seed-docs-ai", version: "1.0.0" },
+    });
+
+    if (initResponse.error) {
+      throw new Error(`MCP initialize failed: ${initResponse.error.message}`);
+    }
+
+    await mcpNotification("notifications/initialized");
+    isInitialized = true;
+  })().finally(() => {
+    initializingPromise = null;
   });
 
-  await mcpRequest("notifications/initialized");
+  return initializingPromise;
 }
 
 /**

docs/lib/ai/tools.ts

@@ -2,7 +2,8 @@ import { tool } from "ai";
 import { z } from "zod";
 
 /**
- * 클라이언트사이드 도구: execute 함수 없음 → useChat에서 UI로 렌더링
+ * 채팅 UI 렌더링용 도구.
+ * 서버에서도 execute를 제공해 tool result가 누락되지 않도록 한다.
  */
 export const clientTools = {
   showComponentExample: tool({
@@ -21,6 +22,7 @@ export const clientTools = {
           'Component example path, e.g., "react/action-button/preview", "react/checkbox/preview"',
         ),
     }),
+    execute: async () => ({ shown: true }),
   }),
 
   showInstallation: tool({
@@ -36,6 +38,7 @@ export const clientTools = {
         )
         .describe('Component name in kebab-case, e.g., "action-button", "checkbox", "tabs"'),
     }),
+    execute: async ({ name }) => ({ shown: true, component: name }),
   }),
 
   showCodeBlock: tool({
@@ -45,5 +48,10 @@ export const clientTools = {
       language: z.string().default("tsx").describe("Programming language"),
       title: z.string().optional().describe("Optional title for the code block"),
     }),
+    execute: async ({ language, title }) => ({
+      shown: true,
+      language,
+      hasTitle: Boolean(title),
+    }),
   }),
 };