ai-chatbot/services/pascal-editor/apps/editor/lib/pascal-agent-mcp.ts

import { existsSync } from 'node:fs'
import path from 'node:path'
import { experimental_createMCPClient, type experimental_MCPClient } from '@ai-sdk/mcp'
import { Experimental_StdioMCPTransport } from '@ai-sdk/mcp/mcp-stdio'
import type { ModelMessage, ToolSet } from 'ai'

const BLOCKED_TOOL_NAMES = new Set([
  'create_project',
  'delete_scene',
  'generate_variants',
  'get_project_status',
  'list_scenes',
  'load_scene',
  'rename_scene',
])

type RequiredTool = ToolSet[string] & {
  execute: NonNullable<ToolSet[string]['execute']>
}

export interface PascalAgentMcpSession {
  close: () => Promise<void>
  saveDraft: (abortSignal?: AbortSignal) => Promise<void>
  sceneName: string
  tools: ToolSet
}

function stringEnvironment(env: NodeJS.ProcessEnv): Record<string, string> {
  return Object.fromEntries(
    Object.entries(env).filter((entry): entry is [string, string] => entry[1] !== undefined),
  )
}

function resolveMcpEntry(env: NodeJS.ProcessEnv): string {
  const configured = env.PASCAL_MCP_ENTRY?.trim()
  if (configured) return path.resolve(configured)

  const candidates = [
    path.resolve(process.cwd(), '../../packages/mcp/dist/bin/pascal-mcp.js'),
    path.resolve(process.cwd(), 'packages/mcp/dist/bin/pascal-mcp.js'),
    path.resolve(process.cwd(), '../packages/mcp/dist/bin/pascal-mcp.js'),
  ]
  const match = candidates.find((candidate) => existsSync(candidate))
  if (!match) {
    throw new Error(
      'Pascal MCP executable was not found. Build @pascal-app/mcp or set PASCAL_MCP_ENTRY.',
    )
  }
  return match
}

function resolveMcpCommand(env: NodeJS.ProcessEnv): string {
  const configured = env.PASCAL_MCP_COMMAND?.trim()
  if (configured) return configured
  return process.versions.bun ? process.execPath : 'bun'
}

function toolExecutionOptions(abortSignal?: AbortSignal) {
  return {
    abortSignal,
    messages: [] as ModelMessage[],
    toolCallId: `pascal-agent-${crypto.randomUUID()}`,
  }
}

function requireTool(tools: ToolSet, name: string): RequiredTool {
  const candidate = tools[name]
  if (!candidate?.execute) {
    throw new Error(`Pascal MCP tool is unavailable: ${name}`)
  }
  return candidate as RequiredTool
}

function withStepTimeout(tool: RequiredTool, timeoutMs: number): RequiredTool {
  return {
    ...tool,
    execute: async (input, options) => {
      const timeoutSignal = AbortSignal.timeout(timeoutMs)
      const abortSignal = options.abortSignal
        ? AbortSignal.any([options.abortSignal, timeoutSignal])
        : timeoutSignal
      return tool.execute(input, { ...options, abortSignal })
    },
  }
}

export function isPascalAgentToolAllowed(name: string): boolean {
  return !BLOCKED_TOOL_NAMES.has(name)
}

function readSceneName(result: unknown, fallback: string): string {
  const structured = (result as { structuredContent?: { name?: unknown } })?.structuredContent
  return typeof structured?.name === 'string' && structured.name.trim()
    ? structured.name.trim()
    : fallback
}

function assertToolSucceeded(result: unknown, operation: string): void {
  if ((result as { isError?: unknown })?.isError === true) {
    const content = (result as { content?: unknown })?.content
    const detail = Array.isArray(content)
      ? content
          .filter((item): item is { text: string; type: 'text' } =>
            Boolean(
              item &&
                typeof item === 'object' &&
                (item as { type?: unknown }).type === 'text' &&
                typeof (item as { text?: unknown }).text === 'string',
            ),
          )
          .map((item) => item.text)
          .join(' ')
          .slice(0, 1000)
      : ''
    throw new Error(`Pascal MCP ${operation} failed${detail ? `: ${detail}` : '.'}`)
  }
}

export async function createPascalAgentMcpSession(params: {
  abortSignal?: AbortSignal
  env?: NodeJS.ProcessEnv
  sceneId: string
  stepTimeoutMs: number
}): Promise<PascalAgentMcpSession> {
  const env = params.env ?? process.env
  const entry = resolveMcpEntry(env)
  const transport = new Experimental_StdioMCPTransport({
    args: [entry, '--stdio'],
    command: resolveMcpCommand(env),
    cwd: path.resolve(path.dirname(entry), '../../../..'),
    env: stringEnvironment(env),
    stderr: 'inherit',
  })

  let client: experimental_MCPClient | null = null
  try {
    client = await experimental_createMCPClient({
      name: 'pascal-editor-agent',
      transport,
      version: '0.1.0',
    })
    const allTools = (await client.tools()) as ToolSet
    const loadScene = requireTool(allTools, 'load_scene')
    const loadResult = await loadScene.execute(
      { id: params.sceneId },
      toolExecutionOptions(params.abortSignal),
    )
    assertToolSucceeded(loadResult, 'load_scene')
    const sceneName = readSceneName(loadResult, params.sceneId)
    const saveScene = withStepTimeout(requireTool(allTools, 'save_scene'), params.stepTimeoutMs)

    const scopedSaveScene: RequiredTool = {
      ...saveScene,
      execute: async (input, options) => {
        const candidate = input && typeof input === 'object' ? input : {}
        const result = await saveScene.execute(
          {
            ...candidate,
            expectedVersion: undefined,
            graph: undefined,
            id: params.sceneId,
            includeCurrentScene: true,
            name: sceneName,
            projectId: undefined,
          },
          options,
        )
        return result
      },
    }

    const tools = Object.fromEntries(
      Object.entries(allTools)
        .filter(([name]) => isPascalAgentToolAllowed(name))
        .map(([name, tool]) => [
          name,
          name === 'save_scene'
            ? scopedSaveScene
            : withStepTimeout(tool as RequiredTool, params.stepTimeoutMs),
        ]),
    ) as ToolSet

    let closed = false
    const close = async () => {
      if (closed) return
      closed = true
      await client?.close()
    }

    return {
      close,
      sceneName,
      tools,
      saveDraft: async (abortSignal) => {
        const result = await scopedSaveScene.execute(
          {
            id: params.sceneId,
            includeCurrentScene: true,
            name: sceneName,
            saveMode: 'draft',
          },
          toolExecutionOptions(abortSignal),
        )
        assertToolSucceeded(result, 'save_scene')
      },
    }
  } catch (error) {
    await client?.close().catch(() => undefined)
    throw error
  }
}