@gemini-designer/mcp-server 0.1.2 → 0.1.29

package/src/generation/litellm-client.ts

@@ -15,107 +15,123 @@ import { Config } from '../config/index.js';
 import type { GeminiUserContent } from './gemini-client.js';
 
 type ChatContentPart =
-    | { type: 'text'; text: string }
-    | { type: 'image_url'; image_url: { url: string } };
+  | { type: 'text'; text: string }
+  | { type: 'image_url'; image_url: { url: string } };
+
+type LiteLLMChatCompletionResponse = {
+  choices?: Array<{
+    message?: {
+      content?: unknown;
+    };
+  }>;
+  usage?: unknown;
+  model?: string;
+};
 
 function normalizeBaseUrl(baseUrl: string): string {
-    const trimmed = baseUrl.replace(/\/+$/, '');
-    // Allow users to provide either http://host:port or http://host:port/v1
-    if (trimmed.endsWith('/v1')) return trimmed;
-    return `${trimmed}/v1`;
+  const trimmed = baseUrl.replace(/\/+$/, '');
+  // Allow users to provide either http://host:port or http://host:port/v1
+  if (trimmed.endsWith('/v1')) return trimmed;
+  return `${trimmed}/v1`;
 }
 
 function inferMimeTypeFromInlineData(mimeType: string): string {
-    // LiteLLM expects a data URL. Ensure mime type is reasonable.
-    if (!mimeType || !mimeType.includes('/')) return 'application/octet-stream';
-    return mimeType;
+  // LiteLLM expects a data URL. Ensure mime type is reasonable.
+  if (!mimeType || !mimeType.includes('/')) return 'application/octet-stream';
+  return mimeType;
 }
 
 function toOpenAIUserContent(userContent: GeminiUserContent): string | ChatContentPart[] {
-    if (typeof userContent === 'string') return userContent;
-
-    const parts: ChatContentPart[] = [];
-    for (const part of userContent) {
-        if ('text' in part) {
-            parts.push({ type: 'text', text: part.text });
-            continue;
-        }
-
-        if ('inlineData' in part) {
-            const mime = inferMimeTypeFromInlineData(part.inlineData.mimeType);
-            const url = `data:${mime};base64,${part.inlineData.data}`;
-            parts.push({ type: 'image_url', image_url: { url } });
-            continue;
-        }
+  if (typeof userContent === 'string') return userContent;
+
+  const parts: ChatContentPart[] = [];
+  for (const part of userContent) {
+    if ('text' in part) {
+      parts.push({ type: 'text', text: part.text });
+      continue;
+    }
+
+    if ('inlineData' in part) {
+      const mime = inferMimeTypeFromInlineData(part.inlineData.mimeType);
+      const url = `data:${mime};base64,${part.inlineData.data}`;
+      parts.push({ type: 'image_url', image_url: { url } });
+      continue;
     }
+  }
 
-    // If the model only received one text part, prefer string (cheaper payload)
-    if (parts.length === 1 && parts[0].type === 'text') return parts[0].text;
-    return parts;
+  // If the model only received one text part, prefer string (cheaper payload)
+  if (parts.length === 1 && parts[0].type === 'text') return parts[0].text;
+  return parts;
 }
 
 export async function generateWithLiteLLM(
-    config: Config,
-    systemPrompt: string,
-    userContent: GeminiUserContent,
-    toolName?: string
+  config: Config,
+  systemPrompt: string,
+  userContent: GeminiUserContent,
+  toolName?: string
 ): Promise<string> {
-    if (!config.litellmEndpoint) {
-        throw new Error('LiteLLM endpoint not configured. Set LITELLM_ENDPOINT.');
-    }
-
-    const baseUrl = normalizeBaseUrl(config.litellmEndpoint);
-    const url = `${baseUrl}/chat/completions`;
-
-    const model = config.litellmModel || config.model;
-
-    const body = {
-        model,
-        messages: [
-            { role: 'system', content: systemPrompt },
-            { role: 'user', content: toOpenAIUserContent(userContent) },
-        ],
-        // Keep default-ish sampling. Let downstream agents control via prompts.
-        temperature: 0.2,
-    };
-
-    const headers: Record<string, string> = {
-        'Content-Type': 'application/json',
+  if (!config.litellmEndpoint) {
+    throw new Error('LiteLLM endpoint not configured. Set LITELLM_ENDPOINT.');
+  }
+
+  const baseUrl = normalizeBaseUrl(config.litellmEndpoint);
+  const url = `${baseUrl}/chat/completions`;
+
+  const model = config.litellmModel || config.model;
+
+  const body = {
+    model,
+    messages: [
+      { role: 'system', content: systemPrompt },
+      { role: 'user', content: toOpenAIUserContent(userContent) },
+    ],
+    // Keep default-ish sampling. Let downstream agents control via prompts.
+    temperature: 0.2,
+  };
+
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/json',
+  };
+
+  if (config.litellmApiKey) {
+    headers.Authorization = `Bearer ${config.litellmApiKey}`;
+  }
+
+  const resp = await fetch(url, {
+    method: 'POST',
+    headers,
+    body: JSON.stringify(body),
+  });
+
+  if (!resp.ok) {
+    const text = await resp.text();
+    throw new Error(`LiteLLM error (${resp.status}): ${text}`);
+  }
+
+  const data = (await resp.json()) as LiteLLMChatCompletionResponse;
+  const content = data.choices?.[0]?.message?.content;
+
+  if (config.debug) {
+    const usage = data.usage;
+    console.error('[litellm] Tool:', toolName || '(unknown)');
+    if (usage) console.error('[litellm] Usage:', usage);
+    console.error('[litellm] Model:', data.model || model);
+  }
+
+  if (typeof content === 'string') return content;
+  // Some providers return arrays; stringify safely.
+  if (Array.isArray(content)) {
+    const partToText = (part: unknown): string => {
+      if (typeof part === 'string') return part;
+      if (part && typeof part === 'object' && 'text' in part) {
+        const text = (part as { text?: unknown }).text;
+        if (typeof text === 'string') return text;
+      }
+      return '';
     };
 
-    if (config.litellmApiKey) {
-        headers.Authorization = `Bearer ${config.litellmApiKey}`;
-    }
-
-    const resp = await fetch(url, {
-        method: 'POST',
-        headers,
-        body: JSON.stringify(body),
-    });
-
-    if (!resp.ok) {
-        const text = await resp.text();
-        throw new Error(`LiteLLM error (${resp.status}): ${text}`);
-    }
-
-    const data = (await resp.json()) as any;
-    const content: unknown = data?.choices?.[0]?.message?.content;
-
-    if (config.debug) {
-        const usage = data?.usage;
-        console.error('[litellm] Tool:', toolName || '(unknown)');
-        if (usage) console.error('[litellm] Usage:', usage);
-        console.error('[litellm] Model:', data?.model || model);
-    }
-
-    if (typeof content === 'string') return content;
-    // Some providers return arrays; stringify safely.
-    if (Array.isArray(content)) {
-        return content
-            .map((p) => (typeof p === 'string' ? p : (p?.text ?? '')))
-            .filter(Boolean)
-            .join('');
-    }
+    return content.map(partToText).filter(Boolean).join('');
+  }
 
-    return '';
+  return '';
 }

package/src/generation/remote-client.ts

@@ -9,95 +9,118 @@ import { Config } from '../config/index.js';
 import type { GeminiUserContent } from './gemini-client.js';
 
 interface GatewayRequest {
-    systemPrompt: string;
-    userPrompt?: string;
-    userParts?: Array<{ text: string } | { inlineData: { mimeType: string; data: string } }>;
-    toolName?: string;
+  systemPrompt: string;
+  userPrompt?: string;
+  userParts?: Array<{ text: string } | { inlineData: { mimeType: string; data: string } }>;
+  toolName?: string;
 }
 
 interface GatewayResponse {
-    content: string;
-    tokensUsed: number;
-    inputTokens?: number;
-    model?: string;
-    costMicrodollars?: number;
-    error?: string;
+  content: string;
+  tokensUsed: number;
+  inputTokens?: number;
+  model?: string;
+  costMicrodollars?: number;
+  error?: string;
 }
 
+export type RemoteQuota = {
+  remainingBudgetUsd: number;
+  maxBudgetUsd: number;
+  spendUsd?: number | null;
+  budgetDuration?: string | null;
+  rpmLimit?: number | null;
+  tpmLimit?: number | null;
+};
+
 /**
  * Generate content via remote gateway
  */
 export async function generateWithRemote(
-    config: Config,
-    systemPrompt: string,
-    userPrompt: GeminiUserContent,
-    toolName?: string
+  config: Config,
+  systemPrompt: string,
+  userPrompt: GeminiUserContent,
+  toolName?: string
 ): Promise<string> {
-    if (!config.remoteEndpoint) {
-        throw new Error('Remote endpoint not configured');
-    }
-
-    if (!config.remoteApiKey) {
-        throw new Error('Remote API key not configured');
-    }
-
-    const request: GatewayRequest = {
-        systemPrompt,
-        toolName,
-    };
-
-    if (typeof userPrompt === 'string') {
-        request.userPrompt = userPrompt;
-    } else {
-        request.userParts = userPrompt as any;
-    }
-
-    const response = await fetch(`${config.remoteEndpoint}/generate`, {
-        method: 'POST',
-        headers: {
-            'Content-Type': 'application/json',
-            Authorization: `Bearer ${config.remoteApiKey}`,
-        },
-        body: JSON.stringify(request),
-    });
-
-    if (!response.ok) {
-        const errorText = await response.text();
-        throw new Error(`Gateway error (${response.status}): ${errorText}`);
-    }
-
-    const data = (await response.json()) as GatewayResponse;
-
-    if (data.error) {
-        throw new Error(`Gateway error: ${data.error}`);
-    }
-
-    if (config.debug) {
-        console.error('[remote] Tokens used:', data.tokensUsed);
-        if (typeof data.inputTokens === 'number') console.error('[remote] Input tokens:', data.inputTokens);
-        if (data.model) console.error('[remote] Model:', data.model);
-    }
-
-    return data.content;
+  if (!config.remoteEndpoint) {
+    throw new Error('Remote endpoint not configured');
+  }
+
+  if (!config.remoteApiKey) {
+    throw new Error('Remote API key not configured');
+  }
+
+  const request: GatewayRequest = {
+    systemPrompt,
+    toolName,
+  };
+
+  if (typeof userPrompt === 'string') {
+    request.userPrompt = userPrompt;
+  } else {
+    request.userParts = userPrompt;
+  }
+
+  const response = await fetch(`${config.remoteEndpoint}/generate`, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${config.remoteApiKey}`,
+    },
+    body: JSON.stringify(request),
+  });
+
+  if (!response.ok) {
+    const errorText = await response.text();
+    throw new Error(`Gateway error (${response.status}): ${errorText}`);
+  }
+
+  const data = (await response.json()) as GatewayResponse;
+
+  if (data.error) {
+    throw new Error(`Gateway error: ${data.error}`);
+  }
+
+  if (config.debug) {
+    console.error('[remote] Tokens used:', data.tokensUsed);
+    if (typeof data.inputTokens === 'number')
+      console.error('[remote] Input tokens:', data.inputTokens);
+    if (data.model) console.error('[remote] Model:', data.model);
+  }
+
+  return data.content;
 }
 
 /**
- * Check remaining quota with the gateway
+ * Check remaining quota/budget with the gateway
  */
-export async function checkQuota(config: Config): Promise<{ remaining: number; limit: number }> {
-    if (!config.remoteEndpoint || !config.remoteApiKey) {
-        throw new Error('Remote configuration missing');
-    }
-
-    const response = await fetch(`${config.remoteEndpoint}/quota`, {
-        headers: {
-            Authorization: `Bearer ${config.remoteApiKey}`,
-        },
-    });
-
-    if (!response.ok) {
-        throw new Error(`Failed to check quota: ${response.status}`);
-    }
-
-    return (await response.json()) as { remaining: number; limit: number };
+export async function checkQuota(config: Config): Promise<RemoteQuota> {
+  if (!config.remoteEndpoint || !config.remoteApiKey) {
+    throw new Error('Remote configuration missing');
+  }
+
+  const response = await fetch(`${config.remoteEndpoint}/quota`, {
+    headers: {
+      Authorization: `Bearer ${config.remoteApiKey}`,
+    },
+  });
+
+  if (!response.ok) {
+    throw new Error(`Failed to check quota: ${response.status}`);
+  }
+
+  const data = (await response.json()) as any;
+
+  if (typeof data?.remainingBudgetUsd !== 'number' || typeof data?.maxBudgetUsd !== 'number') {
+    throw new Error('Invalid quota response from gateway (missing budget fields)');
+  }
+
+  return {
+    remainingBudgetUsd: data.remainingBudgetUsd,
+    maxBudgetUsd: data.maxBudgetUsd,
+    spendUsd: typeof data?.spendUsd === 'number' ? data.spendUsd : null,
+    budgetDuration: typeof data?.budgetDuration === 'string' ? data.budgetDuration : null,
+    rpmLimit: typeof data?.rpmLimit === 'number' ? data.rpmLimit : null,
+    tpmLimit: typeof data?.tpmLimit === 'number' ? data.tpmLimit : null,
+  };
 }

package/src/index.ts

@@ -10,27 +10,27 @@ import { createServer } from './server.js';
 import { loadConfig } from './config/index.js';
 
 async function main() {
-    const config = loadConfig();
+  const config = loadConfig();
 
-    console.error(`[gemini-designer-mcp] Starting in ${config.mode} mode...`);
+  console.error(`[gemini-designer-mcp] Starting in ${config.mode} mode...`);
 
-    const server = await createServer(config);
+  const server = await createServer(config);
 
-    // Handle graceful shutdown
-    process.on('SIGINT', async () => {
-        console.error('[gemini-designer-mcp] Shutting down...');
-        await server.close();
-        process.exit(0);
-    });
+  // Handle graceful shutdown
+  process.on('SIGINT', async () => {
+    console.error('[gemini-designer-mcp] Shutting down...');
+    await server.close();
+    process.exit(0);
+  });
 
-    process.on('SIGTERM', async () => {
-        console.error('[gemini-designer-mcp] Shutting down...');
-        await server.close();
-        process.exit(0);
-    });
+  process.on('SIGTERM', async () => {
+    console.error('[gemini-designer-mcp] Shutting down...');
+    await server.close();
+    process.exit(0);
+  });
 }
 
 main().catch((error) => {
-    console.error('[gemini-designer-mcp] Fatal error:', error);
-    process.exit(1);
+  console.error('[gemini-designer-mcp] Fatal error:', error);
+  process.exit(1);
 });