Linting fixes. Fixed crash when hitting rate limit

2025-06-28 15:47:33 -05:00
parent 75dc51bcb1
commit 10a6502f73
5 changed files with 162 additions and 50 deletions
--- a/eslint.config.ts
+++ b/eslint.config.ts
@@ -9,7 +9,7 @@ export default tseslint.config(
  ...tseslint.configs.strictTypeChecked,
  ...tseslint.configs.stylisticTypeChecked,
  {
-    ignores: ['**/node_modules/*', '**/*.mjs', '**/*.js', 'src/mapper.ts'],
+    ignores: ['**/node_modules/*', '**/*.mjs', '**/*.js'],
  },
  {
    languageOptions: {
--- a/src/chatwrapper.ts
+++ b/src/chatwrapper.ts
@@ -12,6 +12,7 @@ import {
 } from '@google/gemini-cli-core/dist/src/config/models.js';

 import { Content, GeminiResponse, Model } from './types.js';
+import consola from 'consola';

 /* ------------------------------------------------------------------ */
 /* 1.  Build the ContentGenerator exactly like the CLI does           */
@@ -32,6 +33,33 @@ const generatorPromise: Promise<ContentGenerator> = (async () => {
 /* ------------------------------------------------------------------ */
 type GenConfig = Record<string, unknown>;

+const MAX_RETRIES = 3;
+const INITIAL_RETRY_DELAY = 1000; // 1 second
+
+async function withRetry<T>(operation: () => Promise<T>): Promise<T> {
+  let retries = 0;
+  while (true) {
+    try {
+      return await operation();
+    } catch (error) {
+      // Check if it's an Error object with a message property
+      if (!(error instanceof Error) ||
+          !error.message.includes('RESOURCE_EXHAUSTED') ||
+          retries >= MAX_RETRIES) {
+        throw error;
+      }
+      retries++;
+      const delay = INITIAL_RETRY_DELAY * Math.pow(2, retries - 1);
+       
+      consola.error(
+        `Rate limit hit, retrying in ${delay}ms ` +
+        `(attempt ${retries}/${MAX_RETRIES})`,
+      );
+      await new Promise(resolve => setTimeout(resolve, delay));
+    }
+  }
+}
+
 export async function sendChat({
  contents,
  generationConfig = {},
@@ -41,11 +69,11 @@ export async function sendChat({
  tools?: unknown,                // accepted but ignored for now
 }): Promise<GeminiResponse> {
  const generator = await generatorPromise;
-  const gResp = await generator.generateContent({
+  const gResp = await withRetry(() => generator.generateContent({
    model: modelName,
    contents,
    config: generationConfig,
-  });
+  }));
  return {
    text: gResp.text ?? '',
    usageMetadata: {
@@ -65,11 +93,11 @@ export async function* sendChatStream({
  tools?: unknown,
 }) {
  const generator = await generatorPromise;
-  const stream = await generator.generateContentStream({
+  const stream = await withRetry(() => generator.generateContentStream({
    model: modelName,
    contents,
    config: generationConfig,
-  });
+  }));
  for await (const chunk of stream) yield chunk;
 }

@@ -94,3 +122,4 @@ export function listModels(): Model[] {
 // export async function embed(_input: unknown) {
 //   throw new Error('Embeddings endpoint not implemented yet.');
 // }
+
--- a/src/mapper.ts
+++ b/src/mapper.ts
@@ -2,16 +2,25 @@
 /*  mapper.ts – OpenAI ⇆ Gemini (with reasoning/1 M context)           */
 /* ------------------------------------------------------------------ */
 import { fetchAndEncode } from './remoteimage';
-import { z } from 'zod';
-import { ToolRegistry } from '@google/gemini-cli-core/dist/src/tools/tool-registry.js';
-import { RequestBody } from './types';
+import { z, ZodRawShape } from 'zod';
+import { ToolRegistry }
+  from '@google/gemini-cli-core/dist/src/tools/tool-registry.js';
+import { Config } from '@google/gemini-cli-core/dist/src/config/config.js';
+import { Tool } from '@google/gemini-cli-core/dist/src/tools/tools.js';
+import {
+  Part,
+  RequestBody,
+  GeminiResponse,
+  GeminiStreamChunk,
+} from './types';

-/* ------------------------------------------------------------------ */
-interface Part { text?: string; inlineData?: { mimeType: string, data: string } }
-
-/* ------------------------------------------------------------------ */
-function callLocalFunction(_name: string, _args: unknown) {
-  return { ok: true };
+/* ----------------------------------------------------------------- */
+async function callLocalFunction(/*_name: string, _args: unknown*/) {
+  return Promise.resolve({
+    ok: true,
+    llmContent: [],
+    returnDisplay: 'Function executed successfully',
+  });
 }

 /* ================================================================== */
@@ -61,21 +70,29 @@ export async function mapRequest(body: RequestBody) {
  };

  /* ---- Tool / function mapping ----------------------------------- */
-  const tools = new ToolRegistry({} as any);
+  // Note: ToolRegistry expects a complex Config object that we don't have
+  // access to. Casting to `Config` is a workaround.
+  const tools = new ToolRegistry({} as Config);

  if (body.functions?.length) {
-    const reg = tools as any;
-    body.functions.forEach((fn: any) =>
-      reg.registerTool(
-        fn.name,
-        {
-          title: fn.name,
+    for (const fn of body.functions) {
+      tools.registerTool({
+        name: fn.name,
+        displayName: fn.name,
        description: fn.description ?? '',
-          inputSchema: z.object(fn.parameters?.properties ?? {}),
-        },
-        (args: unknown) => callLocalFunction(fn.name, args),
+        schema: z.object(
+          (fn.parameters?.properties as ZodRawShape) ?? {},
        ),
-    );
+        isOutputMarkdown: false,
+        canUpdateOutput: false,
+        validateToolParams: () => null,
+        getDescription: (params: unknown) =>
+          `Executing ${fn.name} with parameters: ` +
+          JSON.stringify(params),
+        shouldConfirmExecute: () => Promise.resolve(false),
+        execute: () => callLocalFunction(),
+      } as Tool);
+    }
  }

  return { geminiReq, tools };
@@ -84,8 +101,12 @@ export async function mapRequest(body: RequestBody) {
 /* ================================================================== */
 /* Non-stream response: Gemini ➞ OpenAI                                */
 /* ================================================================== */
-export function mapResponse(gResp: any) {
-  const usage = gResp.usageMetadata ?? {};
+export function mapResponse(gResp: GeminiResponse) {
+  const usage = gResp.usageMetadata ?? {
+    promptTokens: 0,
+    candidatesTokens: 0,
+    totalTokens: 0,
+  };
  return {
    id: `chatcmpl-${Date.now()}`,
    object: 'chat.completion',
@@ -99,9 +120,9 @@ export function mapResponse(gResp: any) {
      },
    ],
    usage: {
-      prompt_tokens: usage.promptTokens ?? 0,
-      completion_tokens: usage.candidatesTokens ?? 0,
-      total_tokens: usage.totalTokens ?? 0,
+      prompt_tokens: usage.promptTokens,
+      completion_tokens: usage.candidatesTokens,
+      total_tokens: usage.totalTokens,
    },
  };
 }
@@ -110,9 +131,9 @@ export function mapResponse(gResp: any) {
 /* Stream chunk mapper: Gemini ➞ OpenAI                                */
 /* ================================================================== */

-export function mapStreamChunk(chunk: any) {
+export function mapStreamChunk(chunk: GeminiStreamChunk) {
  const part = chunk?.candidates?.[0]?.content?.parts?.[0] ?? {};
-  const delta: any = { role: 'assistant' };
+  const delta: { role: 'assistant', content?: string } = { role: 'assistant' };

  if (part.thought === true) {
    delta.content = `<think>${part.text ?? ''}`;  // ST renders grey bubble
--- a/src/server.ts
+++ b/src/server.ts
@@ -2,7 +2,7 @@ import consola from 'consola';
 import http from 'http';
 import { listModels, sendChat, sendChatStream } from './chatwrapper';
 import { mapRequest, mapResponse, mapStreamChunk } from './mapper.js';
-import { RequestBody, GeminiResponse } from './types';
+import { RequestBody, GeminiResponse, GeminiStreamChunk, Part } from './types';
 import { config } from './config';

 /* ── basic config ─────────────────────────────────────────────────── */
@@ -41,8 +41,11 @@ function readJSON(
              error: { message: 'Request body is missing for POST request' },
            }),
          );
+          resolve(null);
+          return;
        }
-        return resolve(null);
+        resolve(null);
+        return;
      }
      try {
        resolve(JSON.parse(data) as RequestBody);
@@ -51,6 +54,7 @@ function readJSON(
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: { message: 'Malformed JSON' } }));
        resolve(null);
+        return;
      }
    });
  });
@@ -101,7 +105,35 @@ http
          });

          for await (const chunk of sendChatStream({ ...geminiReq, tools })) {
-            res.write(`data: ${JSON.stringify(mapStreamChunk(chunk))}\n\n`);
+            // Transform the chunk to match our expected type
+            const transformedParts = 
+            chunk.candidates?.[0]?.content?.parts?.map(part => {
+              const transformedPart: Part = {
+                text: part.text,
+                thought: part.text?.startsWith?.('<think>') ?? false,
+              };
+              
+              if (part.inlineData?.data) {
+                transformedPart.inlineData = {
+                  mimeType: part.inlineData.mimeType ?? 'text/plain',
+                  data: part.inlineData.data,
+                };
+              }
+              
+              return transformedPart;
+            }) ?? [];
+
+            const streamChunk: GeminiStreamChunk = {
+              candidates: [{
+                content: {
+                  parts: transformedParts,
+                },
+              }],
+            };
+            
+            res.write(
+              `data: ${JSON.stringify(mapStreamChunk(streamChunk))}\n\n`,
+            );
          }
          res.end('data: [DONE]\n\n');
        } else {
@@ -112,10 +144,23 @@ http
      } catch (err) {
        const error = err as Error;
        consola.error('Proxy error ➜', error);
+        
+        // For streaming responses, send error in stream format
+        if (body.stream && res.headersSent) {
+          res.write(`data: ${JSON.stringify({
+            error: {
+              message: error.message,
+              type: 'error',
+            },
+          })}\n\n`);
+          res.end('data: [DONE]\n\n');
+          return;
+        } else {
+          // For non-streaming responses or if headers haven't been sent yet
          res.writeHead(500, { 'Content-Type': 'application/json' });
          res.end(JSON.stringify({ error: { message: error.message } }));
        }
-      return;
+      }
    }

    /* ---- anything else ---------- */
--- a/src/types.ts
+++ b/src/types.ts
@@ -6,9 +6,16 @@ export interface Model {
  object: 'model';
  owned_by: 'google';
 }
+
+interface InlineData {
+   mimeType: string;
+   data: string;
+}
+
 export interface Part {
   text?: string;
-  inlineData?: { mimeType: string, data: string };
+   inlineData?: InlineData;
+   thought?: boolean;
 }

 export interface Content {
@@ -16,6 +23,14 @@ export interface Content {
  parts: Part[];
 }

+interface FunctionDef {
+  name: string;
+  description?: string;
+  parameters?: {
+    properties?: Record<string, unknown>,
+  };
+}
+
 export interface RequestBody {
  messages: {
    content:
@@ -28,13 +43,7 @@ export interface RequestBody {
  generationConfig?: Record<string, unknown>;
  include_reasoning?: boolean;
  stream?: boolean;
-  functions?: {
-    name: string,
-    description?: string,
-    parameters?: {
-      properties?: Record<string, unknown>,
-    },
-  }[];
+  functions?: FunctionDef[];
 }

 export interface GeminiResponse {
@@ -45,3 +54,11 @@ export interface GeminiResponse {
    totalTokens: number,
  };
 }
+
+export interface GeminiStreamChunk {
+  candidates?: {
+    content?: {
+      parts?: Part[],
+    },
+  }[];
+}