aui-mcp-server 0.0.1

package/src/client/types.ts

@@ -0,0 +1,136 @@
+import type { ServerToClientMessage } from '@a2ui/react';
+
+/** A2UI extension URI (from the A2A protocol definition). */
+export const A2UI_EXTENSION_URI = 'https://a2ui.org/a2a-extension/a2ui/v0.8';
+
+/** A2UI DataPart mimeType. */
+export const A2UI_MIME_TYPE = 'application/json+a2ui';
+
+// ── Module Federation x-loader config ───────────────────────────────────────
+
+/**
+ * Runtime config aligned with the `x-loader` field in the MCP Catalog.
+ */
+export interface XLoaderConfig {
+  type: 'module-federation';
+  /** remoteEntry.js URL (passed to MF runtime as the remote entry). */
+  url: string;
+  /** MF remote name (`moduleFederation.name`). */
+  scope: string;
+  /** Exposed module path, e.g. "./ArtistProfile". */
+  module: string;
+}
+
+// ── A2A SSE response types ─────────────────────────────────────────────────-
+
+export interface A2ATaskStatus {
+  state: 'submitted' | 'working' | 'input-required' | 'completed' | 'failed' | 'canceled';
+  /** Inline message from the agent, may contain text/data parts */
+  message?: {
+    kind: 'message';
+    role: 'agent';
+    messageId?: string;
+    contextId?: string;
+    // Narrowed to WirePart[] inside the hook
+    parts: unknown[];
+  };
+  timestamp?: string;
+}
+
+export interface A2AArtifact {
+  artifactId: string;
+  name?: string;
+  parts: unknown[];
+  index?: number;
+  lastChunk?: boolean;
+  metadata?: Record<string, unknown> | null;
+}
+
+/** Task status update event. */
+export interface TaskStatusUpdateEvent {
+  id: string;
+  result: {
+    /** Task ID */
+    id: string;
+    contextId: string;
+    status: A2ATaskStatus;
+    final: boolean;
+    metadata?: Record<string, unknown> | null;
+  };
+}
+
+/** Task artifact update event (usually contains the final A2UI DataParts). */
+export interface TaskArtifactUpdateEvent {
+  id: string;
+  result: {
+    taskId: string;
+    contextId: string;
+    artifact: A2AArtifact;
+  };
+}
+
+export type A2AStreamEvent = TaskStatusUpdateEvent | TaskArtifactUpdateEvent;
+
+// ── Chat / Hook types ─────────────────────────────────────────────────------
+
+export interface ChatMessage {
+  role: 'user' | 'assistant';
+  text: string;
+  /** True when this message triggers an A2UI surface render. */
+  hasSurface?: boolean;
+}
+
+/**
+ * `useA2UIStream` hook options.
+ */
+export interface A2UIStreamOptions {
+  /**
+   * A2A Agent HTTP base URL. Defaults to `/api`.
+   *
+   * - `useA2UIStream({ agentUrl: '/api/a2a' })`
+   */
+  agentUrl?: string;
+
+  /**
+   * Whether `sendMessage` input is treated as an action payload (DataPart) by default.
+   * If true, you can pass a JSON string directly as a `userAction` payload.
+   */
+  defaultIsActionPayload?: boolean;
+}
+
+/**
+ * Logger creation options.
+ */
+export interface CreateLoggerOptions {
+  /** Remote log endpoint. Defaults to `/api/log`. */
+  endpoint?: string;
+  /** Whether to hijack global console.*. Defaults to false. */
+  hijackConsole?: boolean;
+}
+
+/**
+ * MF Registry creation options.
+ */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export interface CreateMFRegistryOptions {
+  /**
+   * Additional shared deps (merged after the default react / react-dom / @a2ui/react config).
+   * Key is the package name; value is a shared config fragment for `@module-federation/enhanced/runtime`.
+   */
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  shared?: Record<string, any>;
+
+  /** Expose `window.__AUI_REGISTRY__` for debugging. Defaults to false. */
+  exposeGlobal?: boolean;
+}
+
+/**
+ * Payload shape for converting tool calls into A2UI messages.
+ */
+export interface ToolCallPayload {
+  tool_name: string;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  tool_input?: Record<string, any>;
+}
+
+export type { ServerToClientMessage };

package/src/client/useA2UIStream.ts

@@ -0,0 +1,378 @@
+import { useState, useCallback, useRef } from 'react';
+import { useA2UIActions } from '@a2ui/react';
+import type { ServerToClientMessage } from '@a2ui/react';
+import {
+  A2UI_EXTENSION_URI,
+  A2UI_MIME_TYPE,
+  type A2AStreamEvent,
+  type TaskStatusUpdateEvent,
+  type TaskArtifactUpdateEvent,
+  type ChatMessage,
+  type A2UIStreamOptions,
+} from './types';
+import { logger } from './useLogger';
+
+// ── A2A wire-format Part types (aligned with the Python SDK) ────────────────
+
+interface WireTextPart {
+  kind: 'text';
+  text: string;
+  metadata?: Record<string, string> | null;
+}
+
+interface WireDataPart {
+  kind: 'data';
+  data: Record<string, unknown>;
+  metadata?: Record<string, string> | null;
+}
+
+type WirePart = WireTextPart | WireDataPart;
+
+// ── Type guards ──────────────────────────────────────────────────────────────
+
+function isTaskStatusUpdateEvent(e: A2AStreamEvent): e is TaskStatusUpdateEvent {
+  return 'result' in e && 'status' in (e as TaskStatusUpdateEvent).result;
+}
+
+function isTaskArtifactUpdateEvent(e: A2AStreamEvent): e is TaskArtifactUpdateEvent {
+  return 'result' in e && 'artifact' in (e as TaskArtifactUpdateEvent).result;
+}
+
+function isA2UIPart(part: WirePart): part is WireDataPart {
+  return (
+    part.kind === 'data' &&
+    (part.metadata?.mimeType === A2UI_MIME_TYPE ||
+      // snake_case compatibility
+      (part.metadata as any)?.mime_type === A2UI_MIME_TYPE)
+  );
+}
+
+// ── Helpers ──────────────────────────────────────────────────────────────────
+
+function extractText(parts: WirePart[]): string {
+  return parts
+    .filter((p): p is WireTextPart => p.kind === 'text')
+    .map((p) => p.text)
+    .join('');
+}
+
+function extractA2UIData(parts: WirePart[]): Record<string, unknown>[] {
+  return parts.filter(isA2UIPart).map((p) => p.data);
+}
+
+// ── Hook ─────────────────────────────────────────────────────────────────────
+
+/**
+ * A2A JSON-RPC 2.0 + SSE streaming hook.
+ *
+ * - Defaults to `/api` as the Agent HTTP entry
+ * - Supports `isActionPayload` mode to serialize user input as a DataPart
+ */
+export interface A2UIStreamPerfSnapshot {
+  /** t0: request start timestamp (message/stream), in ms (performance.now) */
+  t0: number;
+  /** t1: timestamp when the first A2UI DataPart triggers render, in ms (performance.now) */
+  t1?: number;
+  /** t2: completion timestamp (SSE end or final status), in ms (performance.now) */
+  t2?: number;
+  /** TTFC = t1 - t0, in ms; undefined if no A2UI DataPart appears in this round */
+  ttfcMs?: number;
+  /** TTLC = t2 - t0, in ms */
+  ttlcMs?: number;
+  /** Completion reason for debugging (e.g. 'status-completed' | 'sse-end') */
+  reason?: string;
+}
+
+export function useA2UIStream(options?: A2UIStreamOptions) {
+  const actions = useA2UIActions();
+
+  const agentUrl = options?.agentUrl ?? '/api';
+  const defaultIsActionPayload = options?.defaultIsActionPayload ?? false;
+
+  const [messages, setMessages] = useState<ChatMessage[]>([]);
+  const [isLoading, setIsLoading] = useState(false);
+  const [lastPerf, setLastPerf] = useState<A2UIStreamPerfSnapshot | null>(null);
+
+  // `contextId` stays stable across the session; `taskId` persists after the first response
+  const contextIdRef = useRef<string>(crypto.randomUUID());
+  const taskIdRef = useRef<string | undefined>(undefined);
+
+  // Perf timing state for the current round (use performance.now to avoid Date precision issues)
+  const perfStateRef = useRef<{
+    t0?: number;
+    t1?: number;
+    t2?: number;
+    hasFirstComponent?: boolean;
+    completed?: boolean;
+  }>({});
+
+  const now = () => (typeof performance !== 'undefined' && performance.now ? performance.now() : Date.now());
+
+  const markFirstComponentRendered = () => {
+    const state = perfStateRef.current;
+    if (!state.t0 || state.t1 || state.completed) return;
+    state.t1 = now();
+  };
+
+  const finalizePerf = (reason: string) => {
+    const state = perfStateRef.current;
+    if (!state.t0 || state.completed) return;
+
+    state.completed = true;
+    state.t2 = state.t2 ?? now();
+
+    const t0 = state.t0;
+    const t1 = state.t1;
+    const t2 = state.t2;
+    const ttfcMs = t1 ? t1 - t0 : undefined;
+    const ttlcMs = t2 - t0;
+
+    const snapshot: A2UIStreamPerfSnapshot = {
+      t0,
+      t1,
+      t2,
+      ttfcMs,
+      ttlcMs,
+      reason,
+    };
+
+    setLastPerf(snapshot);
+
+    const ttfcLabel = ttfcMs !== undefined ? `${(ttfcMs / 1000).toFixed(2)}s` : 'N/A';
+    const totalLabel = `${(ttlcMs / 1000).toFixed(2)}s`;
+
+    // Print a visible perf snapshot to the console for quick inspection
+    // eslint-disable-next-line no-console
+    console.log('[AUI-X Perf] TTFC: %s, Total: %s', ttfcLabel, totalLabel, {
+      t0,
+      t1,
+      t2,
+      ttfcMs,
+      ttlcMs,
+      reason,
+    });
+  };
+
+  const sendMessage = useCallback(
+    async (userText: string, isActionPayload = defaultIsActionPayload) => {
+      if (!userText.trim() || isLoading) return;
+
+      setIsLoading(true);
+      // Reset timers per round and record t0
+      perfStateRef.current = {
+        t0: now(),
+        t1: undefined,
+        t2: undefined,
+        hasFirstComponent: false,
+        completed: false,
+      };
+      if (!isActionPayload) {
+        setMessages((prev) => [...prev, { role: 'user', text: userText }]);
+      }
+
+      let assistantText = '';
+      let renderedSurface = false;
+
+      try {
+        // ── Build message.parts ─────────────────────────────────────────
+        let parts: WirePart[];
+        if (isActionPayload) {
+          let actionData: Record<string, unknown>;
+          try {
+            actionData = JSON.parse(userText) as Record<string, unknown>;
+          } catch {
+            actionData = { userAction: { actionName: userText, context: {} } };
+          }
+          parts = [{ kind: 'data', data: actionData }];
+        } else {
+          parts = [{ kind: 'text', text: userText }];
+        }
+
+        // ── Build A2A JSON-RPC request ──────────────────────────────────
+        const rpcRequest = {
+          jsonrpc: '2.0' as const,
+          id: crypto.randomUUID(),
+          method: 'message/stream' as const,
+          params: {
+            message: {
+              kind: 'message' as const,
+              role: 'user' as const,
+              messageId: crypto.randomUUID(),
+              contextId: contextIdRef.current,
+              ...(taskIdRef.current ? { taskId: taskIdRef.current } : {}),
+              extensions: [A2UI_EXTENSION_URI],
+              parts,
+            },
+          },
+        };
+
+        const response = await fetch(agentUrl, {
+          method: 'POST',
+          headers: { 'Content-Type': 'application/json' },
+          body: JSON.stringify(rpcRequest),
+        });
+
+        if (!response.ok) {
+          throw new Error(`HTTP ${response.status}: ${response.statusText}`);
+        }
+        if (!response.body) {
+          throw new Error('No response body');
+        }
+
+        // ── Parse SSE stream ───────────────────────────────────────────
+        const reader = response.body.getReader();
+        const decoder = new TextDecoder();
+        let buffer = '';
+
+        // eslint-disable-next-line no-constant-condition
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+
+          buffer += decoder.decode(value, { stream: true });
+          const lines = buffer.split('\n');
+          buffer = lines.pop() ?? '';
+
+          for (const line of lines) {
+            if (!line.startsWith('data: ')) continue;
+            const jsonStr = line.slice(6).trim();
+            if (!jsonStr || jsonStr === '[DONE]') continue;
+
+            let event: A2AStreamEvent;
+            try {
+              event = JSON.parse(jsonStr) as A2AStreamEvent;
+            } catch {
+              continue;
+            }
+
+            // ── TaskStatusUpdateEvent ─────────────────────────────────
+            if (isTaskStatusUpdateEvent(event)) {
+              const { id: taskId, status } = event.result;
+
+              if (taskId) taskIdRef.current = taskId;
+
+              const rawParts = (status.message?.parts ?? []) as WirePart[];
+              const textChunk = extractText(rawParts);
+              if (textChunk) {
+                assistantText += textChunk;
+                setMessages((prev) => {
+                  const last = prev[prev.length - 1];
+                  if (last?.role === 'assistant') {
+                    return [...prev.slice(0, -1), { ...last, text: assistantText }];
+                  }
+                  return [...prev, { role: 'assistant', text: assistantText }];
+                });
+              }
+
+              const a2uiData = extractA2UIData(rawParts);
+              if (a2uiData.length > 0) {
+                // Record t1 (TTFC) when the first A2UI DataPart triggers render
+                markFirstComponentRendered();
+                logger.info('[A2UI] Received UI update part', {
+                  taskId,
+                  componentCount: a2uiData.length,
+                });
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                actions.processMessages(a2uiData as any);
+                const surfaces = actions.getSurfaces();
+                logger.debug('[A2UI] Surfaces after process:', Array.from(surfaces.keys()));
+
+                renderedSurface = true;
+                setMessages((prev) => {
+                  const last = prev[prev.length - 1];
+                  if (last?.role === 'assistant') {
+                    return [...prev.slice(0, -1), { ...last, hasSurface: true }];
+                  }
+                  return [...prev, { role: 'assistant', text: assistantText, hasSurface: true }];
+                });
+              }
+            }
+
+            // ── TaskArtifactUpdateEvent ───────────────────────────────
+            if (isTaskArtifactUpdateEvent(event)) {
+              const { artifact } = event.result;
+              const rawParts = (artifact.parts ?? []) as WirePart[];
+
+              const textChunk = extractText(rawParts);
+              if (textChunk) {
+                assistantText += textChunk;
+                setMessages((prev) => {
+                  const last = prev[prev.length - 1];
+                  if (last?.role === 'assistant') {
+                    return [...prev.slice(0, -1), { ...last, text: assistantText }];
+                  }
+                  return [...prev, { role: 'assistant', text: assistantText }];
+                });
+              }
+
+              const a2uiData = extractA2UIData(rawParts);
+              if (a2uiData.length > 0) {
+                // Record t1 (TTFC) when the first A2UI DataPart triggers render
+                markFirstComponentRendered();
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                actions.processMessages(a2uiData as any);
+                renderedSurface = true;
+                setMessages((prev) => {
+                  const last = prev[prev.length - 1];
+                  if (last?.role === 'assistant') {
+                    return [...prev.slice(0, -1), { ...last, hasSurface: true }];
+                  }
+                  return [...prev, { role: 'assistant', text: assistantText, hasSurface: true }];
+                });
+              }
+            }
+
+            // If the event is final, consider this round completed and record t2 (TTLC)
+            if (
+              isTaskStatusUpdateEvent(event) &&
+              (event.result.final || event.result.status.state === 'completed')
+            ) {
+              finalizePerf('status-completed');
+            }
+          }
+        }
+
+        // Fallback: ensure at least one assistant bubble
+        if (!assistantText && !renderedSurface) {
+          setMessages((prev) => [...prev, { role: 'assistant', text: '…' }]);
+        }
+      } catch (err) {
+        setMessages((prev) => [
+          ...prev,
+          { role: 'assistant', text: `Connection error: ${String(err)}` },
+        ]);
+      } finally {
+        setIsLoading(false);
+        // Even if we didn't receive a final status, record t2 when SSE ends
+        finalizePerf('sse-end');
+      }
+    },
+    [agentUrl, defaultIsActionPayload, isLoading, actions],
+  );
+
+  const processLocalMessages = useCallback(
+    (serverMessages: ServerToClientMessage[], label = 'Local tool call: rendered UI surface') => {
+      if (!serverMessages || serverMessages.length === 0) return;
+
+      logger.info('[A2UI] processLocalMessages:', label, `(${serverMessages.length} messages)`);
+      logger.debug('[A2UI] Surfaces BEFORE process:', Array.from(actions.getSurfaces().keys()));
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      actions.processMessages(serverMessages as any);
+
+      logger.debug('[A2UI] Surfaces AFTER process:', Array.from(actions.getSurfaces().keys()));
+
+      setMessages((prev) => [
+        ...prev,
+        {
+          role: 'assistant',
+          text: label,
+          hasSurface: true,
+        },
+      ]);
+    },
+    [actions],
+  );
+
+  return { messages, isLoading, sendMessage, processLocalMessages, lastPerf };
+}

package/src/client/useLogger.ts

@@ -0,0 +1,147 @@
+import { useEffect, useRef } from 'react';
+import type { CreateLoggerOptions } from './types';
+
+type LogLevel = 'DEBUG' | 'INFO' | 'WARN' | 'ERROR';
+
+class RemoteLogger {
+  private queue: unknown[] = [];
+  private isProcessing = false;
+
+  constructor(private endpoint: string) {}
+
+  setEndpoint(endpoint: string) {
+    this.endpoint = endpoint;
+  }
+
+  private async flush() {
+    if (this.isProcessing || this.queue.length === 0) return;
+    this.isProcessing = true;
+
+    while (this.queue.length > 0) {
+      const entry = this.queue.shift();
+      try {
+        await fetch(this.endpoint, {
+          method: 'POST',
+          headers: { 'Content-Type': 'application/json' },
+          body: JSON.stringify(entry),
+        });
+      } catch (e) {
+        // Use the original console to avoid recursion
+        // eslint-disable-next-line no-console
+        console.warn('[AUI Logger] Failed to sync log to backend', e);
+        // Abort the current flush to avoid an infinite loop when backend is unavailable
+        break;
+      }
+    }
+
+    this.isProcessing = false;
+  }
+
+  log(level: LogLevel, ...args: unknown[]) {
+    const time = new Date().toISOString();
+    const message = args
+      .map((arg) => (typeof arg === 'object' ? JSON.stringify(arg, null, 2) : String(arg)))
+      .join(' ');
+
+    const consoleMethod = level.toLowerCase() as 'debug' | 'info' | 'warn' | 'error';
+
+    // Prefer the preserved original console methods (if present)
+    const original =
+      (typeof window !== 'undefined' && (window as any).__AUI_ORIGINAL_CONSOLE__?.[consoleMethod]) ||
+      // eslint-disable-next-line no-console
+      console[consoleMethod];
+
+    original(`[${level}]`, ...args);
+
+    this.queue.push({ time, level, message });
+    void this.flush();
+  }
+}
+
+export interface Logger {
+  debug: (...args: unknown[]) => void;
+  info: (...args: unknown[]) => void;
+  warn: (...args: unknown[]) => void;
+  error: (...args: unknown[]) => void;
+}
+
+let remoteLogger: RemoteLogger | null = null;
+let globalLogger: Logger | null = null;
+
+function ensureRemoteLogger(endpoint: string): RemoteLogger {
+  if (!remoteLogger) {
+    remoteLogger = new RemoteLogger(endpoint);
+  } else {
+    remoteLogger.setEndpoint(endpoint);
+  }
+  return remoteLogger;
+}
+
+function hijackConsoleIfNeeded(logger: Logger) {
+  if (typeof window === 'undefined') return;
+
+  const w = window as any;
+  if (w.__AUI_CONSOLE_HIJACKED__) return;
+
+  w.__AUI_ORIGINAL_CONSOLE__ = {
+    debug: console.debug,
+    info: console.info,
+    log: console.log,
+    warn: console.warn,
+    error: console.error,
+  };
+
+  console.debug = (...args) => logger.debug(...args);
+  console.info = (...args) => logger.info(...args);
+  console.log = (...args) => logger.info(...args);
+  console.warn = (...args) => logger.warn(...args);
+  console.error = (...args) => logger.error(...args);
+
+  w.__AUI_CONSOLE_HIJACKED__ = true;
+}
+
+/**
+ * Create (or reconfigure) the global logger.
+ *
+ * - `endpoint`: remote log endpoint (default: `/api/log`)
+ * - `hijackConsole`: whether to hijack console.* (default: false)
+ */
+export function createLogger(options?: CreateLoggerOptions): Logger {
+  const endpoint = options?.endpoint ?? '/api/log';
+  const hijackConsole = options?.hijackConsole ?? false;
+
+  const remote = ensureRemoteLogger(endpoint);
+  const logger: Logger = {
+    debug: (...args: unknown[]) => remote.log('DEBUG', ...args),
+    info: (...args: unknown[]) => remote.log('INFO', ...args),
+    warn: (...args: unknown[]) => remote.log('WARN', ...args),
+    error: (...args: unknown[]) => remote.log('ERROR', ...args),
+  };
+
+  if (hijackConsole) {
+    hijackConsoleIfNeeded(logger);
+  }
+
+  globalLogger = logger;
+  return logger;
+}
+
+/**
+ * Default exported logger:
+ * - endpoint: `/api/log`
+ * - does not hijack console (configure on demand)
+ */
+export const logger: Logger = createLogger();
+
+/**
+ * React hook: initialize the logger on mount and optionally hijack console.*
+ */
+export function useLogger(options?: CreateLoggerOptions): Logger {
+  const optionsRef = useRef(options);
+
+  useEffect(() => {
+    createLogger(optionsRef.current);
+  }, []);
+
+  return globalLogger ?? logger;
+}