feat: migrate gateway to nestjs

src/server/modules/proxy/clients/gemini.client.ts

@@ -0,0 +1,59 @@
+import { Injectable, Logger } from '@nestjs/common';
+import axios from 'axios';
+import { GeminiRequest, GeminiResponse } from '../interfaces/request-interfaces';
+
+@Injectable()
+export class GeminiClient {
+  private readonly logger = new Logger(GeminiClient.name);
+  // Default to v1beta for most features
+  private readonly baseUrl = 'https://generativelanguage.googleapis.com/v1beta';
+
+  async streamGenerate(model: string, content: GeminiRequest, accessToken: string): Promise<any> {
+    const url = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`;
+
+    try {
+      const response = await axios.post(url, content, {
+        headers: {
+          Authorization: `Bearer ${accessToken}`,
+          'Content-Type': 'application/json',
+        },
+        responseType: 'stream',
+        timeout: 60000,
+      });
+      return response.data;
+    } catch (error) {
+      if (axios.isAxiosError(error)) {
+        this.logger.error(`Gemini Stream API Error: ${error.message}`);
+        throw new Error(error.response?.data?.error?.message || error.message);
+      }
+      throw error;
+    }
+  }
+
+  async generate(
+    model: string,
+    content: GeminiRequest,
+    accessToken: string,
+  ): Promise<GeminiResponse> {
+    const url = `${this.baseUrl}/models/${model}:generateContent`;
+
+    try {
+      const response = await axios.post<GeminiResponse>(url, content, {
+        headers: {
+          Authorization: `Bearer ${accessToken}`,
+          'Content-Type': 'application/json',
+        },
+        timeout: 60000, // 60s timeout
+      });
+      return response.data;
+    } catch (error) {
+      if (axios.isAxiosError(error)) {
+        this.logger.error(
+          `Gemini API Error: ${error.message} - ${JSON.stringify(error.response?.data)}`,
+        );
+        throw new Error(error.response?.data?.error?.message || error.message);
+      }
+      throw error;
+    }
+  }
+}

src/server/modules/proxy/interfaces/request-interfaces.ts

@@ -0,0 +1,141 @@
+export interface OpenAIChatRequest {
+  model: string;
+  messages: OpenAIMessage[];
+  temperature?: number;
+  top_p?: number;
+  max_tokens?: number;
+  stream?: boolean;
+  size?: string;
+  quality?: string;
+  extra?: Record<string, any>;
+}
+
+export interface OpenAIMessage {
+  role: string;
+  content: string | OpenAIContentPart[];
+  name?: string;
+}
+
+export interface OpenAIContentPart {
+  type: 'text' | 'image_url';
+  text?: string;
+  image_url?: {
+    url: string;
+  };
+}
+
+export interface AnthropicChatRequest {
+  model: string;
+  messages: AnthropicMessage[];
+  system?: string | AnthropicSystemBlock[];
+  max_tokens?: number;
+  metadata?: any;
+  stop_sequences?: string[];
+  stream?: boolean;
+  temperature?: number;
+  top_p?: number;
+  top_k?: number;
+}
+
+export interface AnthropicMessage {
+  role: string;
+  content: string | AnthropicContent[];
+}
+
+export interface AnthropicSystemBlock {
+  type: string;
+  text: string;
+}
+
+export type AnthropicContent =
+  | { type: 'text'; text: string }
+  | { type: 'thinking'; thinking: string; signature?: string }
+  | { type: 'image'; source: AnthropicImageSource };
+
+export interface AnthropicImageSource {
+  type: 'base64';
+  media_type: string;
+  data: string;
+}
+
+export interface GeminiContent {
+  role: string;
+  parts: GeminiPart[];
+}
+
+export interface GeminiPart {
+  text?: string;
+  inlineData?: GeminiInlineData;
+  thoughtSignature?: string;
+}
+
+export interface GeminiInlineData {
+  mimeType: string;
+  data: string;
+}
+
+export interface GeminiRequest {
+  contents: GeminiContent[];
+  systemInstruction?: { parts: GeminiPart[] };
+  generationConfig?: GeminiGenerationConfig;
+}
+
+export interface GeminiGenerationConfig {
+  temperature?: number;
+  maxOutputTokens?: number;
+  topP?: number;
+  topK?: number;
+}
+
+export interface GeminiResponse {
+  candidates?: GeminiCandidate[];
+  usageMetadata?: GeminiUsageMetadata;
+}
+
+export interface GeminiCandidate {
+  content?: GeminiContent;
+  finishReason?: string;
+  index?: number;
+}
+
+export interface GeminiUsageMetadata {
+  promptTokenCount?: number;
+  candidatesTokenCount?: number;
+  totalTokenCount?: number;
+}
+
+export interface OpenAIChatResponse {
+  id: string;
+  object: string;
+  created: number;
+  model: string;
+  choices: OpenAIChoice[];
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
+}
+
+export interface OpenAIChoice {
+  index: number;
+  message: {
+    role: string;
+    content: string;
+  };
+  finish_reason: string;
+}
+
+export interface AnthropicChatResponse {
+  id: string;
+  type: string;
+  role: string;
+  model: string;
+  content: AnthropicContent[];
+  stop_reason: string | null;
+  stop_sequence: string | null;
+  usage: {
+    input_tokens: number;
+    output_tokens: number;
+  };
+}

src/server/modules/proxy/proxy.controller.ts

@@ -0,0 +1,90 @@
+import { Controller, Get, Post, Body, Res, HttpStatus } from '@nestjs/common';
+import { FastifyReply } from 'fastify';
+import { ProxyService } from './proxy.service';
+import { Observable } from 'rxjs';
+import { OpenAIChatRequest, AnthropicChatRequest } from './interfaces/request-interfaces';
+
+@Controller('v1')
+export class ProxyController {
+  constructor(private readonly proxyService: ProxyService) {}
+
+  @Get('models')
+  getModels() {
+    return {
+      object: 'list',
+      data: [
+        // Gemini Native
+        {
+          id: 'gemini-2.5-flash-thinking',
+          object: 'model',
+          created: 1734336000,
+          owned_by: 'google',
+        },
+        { id: 'gemini-2.5-flash', object: 'model', created: 1734336000, owned_by: 'google' },
+        { id: 'gemini-2.5-pro', object: 'model', created: 1734336000, owned_by: 'google' },
+        { id: 'gemini-3-flash', object: 'model', created: 1734336000, owned_by: 'google' },
+
+        // Claude Native (Mapped)
+        { id: 'claude-sonnet-4-5', object: 'model', created: 1734336000, owned_by: 'anthropic' },
+        {
+          id: 'claude-sonnet-4-5-thinking',
+          object: 'model',
+          created: 1734336000,
+          owned_by: 'anthropic',
+        },
+
+        // Image Models
+        { id: 'gemini-3-pro-image', object: 'model', created: 1734336000, owned_by: 'google' },
+      ],
+    };
+  }
+
+  @Post('chat/completions')
+  async chatCompletions(@Body() body: OpenAIChatRequest, @Res() res: FastifyReply) {
+    try {
+      const result = await this.proxyService.handleChatCompletions(body);
+
+      if (body.stream && result instanceof Observable) {
+        res.header('Content-Type', 'text/event-stream');
+        res.header('Cache-Control', 'no-cache');
+        res.header('Connection', 'keep-alive');
+        res.send(result);
+      } else {
+        res.status(HttpStatus.OK).send(result);
+      }
+    } catch (error) {
+      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      res.status(HttpStatus.INTERNAL_SERVER_ERROR).send({
+        error: {
+          message: message,
+          type: 'server_error',
+        },
+      });
+    }
+  }
+
+  @Post('messages')
+  async anthropicMessages(@Body() body: AnthropicChatRequest, @Res() res: FastifyReply) {
+    try {
+      const result = await this.proxyService.handleAnthropicMessages(body);
+
+      if (body.stream && result instanceof Observable) {
+        res.header('Content-Type', 'text/event-stream');
+        res.header('Cache-Control', 'no-cache');
+        res.header('Connection', 'keep-alive');
+        res.send(result);
+      } else {
+        res.status(HttpStatus.OK).send(result);
+      }
+    } catch (error) {
+      const message = error instanceof Error ? error.message : 'Internal Server Error';
+      res.status(HttpStatus.INTERNAL_SERVER_ERROR).send({
+        type: 'error',
+        error: {
+          type: 'api_error',
+          message: message,
+        },
+      });
+    }
+  }
+}

src/server/modules/proxy/proxy.module.ts

@@ -0,0 +1,11 @@
+import { Module } from '@nestjs/common';
+import { ProxyController } from './proxy.controller';
+import { ProxyService } from './proxy.service';
+import { TokenManagerService } from './token-manager.service';
+import { GeminiClient } from './clients/gemini.client';
+
+@Module({
+  controllers: [ProxyController],
+  providers: [ProxyService, TokenManagerService, GeminiClient],
+})
+export class ProxyModule {}

src/server/modules/proxy/proxy.service.ts

@@ -0,0 +1,506 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { TokenManagerService } from './token-manager.service';
+import { GeminiClient } from './clients/gemini.client';
+import { v4 as uuidv4 } from 'uuid';
+import { Observable } from 'rxjs';
+import {
+  OpenAIChatRequest,
+  AnthropicChatRequest,
+  GeminiContent,
+  GeminiPart,
+  AnthropicSystemBlock,
+  OpenAIContentPart,
+  GeminiResponse,
+  GeminiRequest,
+  AnthropicChatResponse,
+  OpenAIChatResponse,
+  GeminiCandidate,
+} from './interfaces/request-interfaces';
+
+@Injectable()
+export class ProxyService {
+  private readonly logger = new Logger(ProxyService.name);
+
+  constructor(
+    private readonly tokenManager: TokenManagerService,
+    private readonly geminiClient: GeminiClient,
+  ) {}
+
+  // --- Anthropic Handlers ---
+
+  async handleAnthropicMessages(
+    request: AnthropicChatRequest,
+  ): Promise<AnthropicChatResponse | Observable<string>> {
+    const targetModel = this.mapModel(request.model);
+    this.logger.log(
+      `Received Anthropic request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`,
+    );
+
+    // Retry loop
+    let lastError: unknown = null;
+    const maxRetries = 3;
+
+    for (let i = 0; i < maxRetries; i++) {
+      const token = await this.tokenManager.getNextToken();
+      if (!token) {
+        throw new Error('No available accounts');
+      }
+
+      try {
+        // Convert Anthropic -> Gemini
+        const geminiRequest = this.convertAnthropicToGemini(request);
+
+        if (request.stream) {
+          const stream = await this.geminiClient.streamGenerate(
+            targetModel,
+            geminiRequest,
+            token.token.access_token,
+          );
+          return this.processAnthropicStreamResponse(stream, request.model);
+        } else {
+          const response = await this.geminiClient.generate(
+            targetModel,
+            geminiRequest,
+            token.token.access_token,
+          );
+          return this.convertGeminiToAnthropicResponse(response, request.model);
+        }
+      } catch (error) {
+        lastError = error;
+        if (error instanceof Error) {
+          this.logger.warn(`Anthropic Request failed: ${error.message}`);
+          if (this.shouldRetry(error.message)) {
+            this.tokenManager.markAsRateLimited(token.email);
+          }
+        }
+      }
+    }
+    throw lastError || new Error('Request failed after retries');
+  }
+
+  private convertAnthropicToGemini(request: AnthropicChatRequest): GeminiRequest {
+    const contents: GeminiContent[] = [];
+    let systemInstruction: { parts: GeminiPart[] } | undefined = undefined;
+
+    if (request.system) {
+      let systemText = '';
+      if (typeof request.system === 'string') {
+        systemText = request.system;
+      } else if (Array.isArray(request.system)) {
+        systemText = request.system
+          .filter((b: AnthropicSystemBlock) => b.type === 'text')
+          .map((b: AnthropicSystemBlock) => b.text)
+          .join('\n');
+      }
+
+      systemInstruction = {
+        parts: [{ text: systemText }],
+      };
+    }
+
+    // Anthropic 'messages' -> Gemini 'contents'
+    for (const msg of request.messages) {
+      const role = msg.role === 'user' ? 'user' : 'model';
+      // Handle content arrays or strings
+      let textPart = '';
+      if (Array.isArray(msg.content)) {
+        textPart = msg.content
+          .filter((c) => c.type === 'text')
+          .map((c) => (c as { type: 'text'; text: string }).text)
+          .join('\n');
+      } else {
+        textPart = msg.content;
+      }
+
+      contents.push({
+        role,
+        parts: [{ text: textPart }],
+      });
+    }
+
+    return {
+      contents,
+      systemInstruction,
+      generationConfig: {
+        temperature: request.temperature,
+        maxOutputTokens: request.max_tokens,
+        topP: request.top_p,
+        topK: request.top_k,
+      },
+      // Note: Rust version handles 'pending_images' logic here, but for now we assume simple text/image structure.
+      // If we need to support image injection from history to user msg, that logic needs to be ported too.
+      // Current port matches previous 'any' style logic but with types.
+    };
+  }
+
+  private convertGeminiToAnthropicResponse(
+    geminiResponse: GeminiResponse,
+    model: string,
+  ): AnthropicChatResponse {
+    const candidate = geminiResponse.candidates?.[0];
+    const processed = this.processInlineData(candidate);
+    const parts = processed?.content?.parts || [];
+    const contentText = parts.map((p) => p.text || '').join('');
+
+    return {
+      id: `msg_${uuidv4()}`,
+      type: 'message',
+      role: 'assistant',
+      model: model,
+      content: [{ type: 'text', text: contentText }],
+      stop_reason: 'end_turn',
+      stop_sequence: null,
+      usage: {
+        input_tokens: geminiResponse.usageMetadata?.promptTokenCount || 0,
+        output_tokens: geminiResponse.usageMetadata?.candidatesTokenCount || 0,
+      },
+    };
+  }
+
+  private processAnthropicStreamResponse(upstreamStream: any, model: string): Observable<string> {
+    return new Observable<string>((subscriber) => {
+      const decoder = new TextDecoder();
+      let buffer = '';
+      const msgId = `msg_${uuidv4()}`;
+
+      // Send Message Start
+      subscriber.next(
+        `event: message_start\ndata: ${JSON.stringify({
+          type: 'message_start',
+          message: {
+            id: msgId,
+            type: 'message',
+            role: 'assistant',
+            content: [],
+            model: model,
+            stop_reason: null,
+            stop_sequence: null,
+            usage: { input_tokens: 0, output_tokens: 0 },
+          },
+        })}\n\n`,
+      );
+
+      subscriber.next(
+        `event: content_block_start\ndata: ${JSON.stringify({
+          type: 'content_block_start',
+          index: 0,
+          content_block: { type: 'text', text: '' },
+        })}\n\n`,
+      );
+
+      upstreamStream.on('data', (chunk: Buffer) => {
+        buffer += decoder.decode(chunk, { stream: true });
+        const lines = buffer.split('\n');
+        buffer = lines.pop() || '';
+
+        for (const line of lines) {
+          const trimmed = line.trim();
+          if (!trimmed.startsWith('data: ')) continue;
+          const dataStr = trimmed.slice(6);
+          if (dataStr === '[DONE]') continue;
+
+          try {
+            const json = JSON.parse(dataStr);
+            const candidate = json.candidates?.[0];
+            const contentPart = candidate?.content?.parts?.[0];
+            const text = contentPart?.text || '';
+
+            if (text) {
+              subscriber.next(
+                `event: content_block_delta\ndata: ${JSON.stringify({
+                  type: 'content_block_delta',
+                  index: 0,
+                  delta: { type: 'text_delta', text: text },
+                })}\n\n`,
+              );
+            }
+          } catch (e) {}
+        }
+      });
+
+      upstreamStream.on('end', () => {
+        subscriber.next(
+          `event: content_block_stop\ndata: ${JSON.stringify({
+            type: 'content_block_stop',
+            index: 0,
+          })}\n\n`,
+        );
+
+        subscriber.next(
+          `event: message_delta\ndata: ${JSON.stringify({
+            type: 'message_delta',
+            delta: { stop_reason: 'end_turn', stop_sequence: null },
+            usage: { output_tokens: 0 },
+          })}\n\n`,
+        );
+
+        subscriber.next(
+          `event: message_stop\ndata: ${JSON.stringify({
+            type: 'message_stop',
+          })}\n\n`,
+        );
+
+        subscriber.complete();
+      });
+
+      upstreamStream.on('error', (err: any) => subscriber.error(err));
+    });
+  }
+
+  // --- OpenAI / Universal Handlers ---
+
+  async handleChatCompletions(
+    request: OpenAIChatRequest,
+  ): Promise<OpenAIChatResponse | Observable<string>> {
+    const targetModel = this.mapModel(request.model);
+    this.logger.log(
+      `Received request for model: ${request.model} (Mapped: ${targetModel}, Stream: ${request.stream})`,
+    );
+
+    // Retry loop for account selection
+    let lastError: unknown = null;
+    const maxRetries = 3;
+
+    for (let i = 0; i < maxRetries; i++) {
+      // 1. Get Token
+      const token = await this.tokenManager.getNextToken();
+      if (!token) {
+        throw new Error('No available accounts (all exhausted or rate limited)');
+      }
+
+      try {
+        // 2. Convert Request
+        const geminiRequest = this.convertOpenAIToGemini(request);
+
+        // 3. Handle Stream vs Non-Stream
+        if (request.stream) {
+          const stream = await this.geminiClient.streamGenerate(
+            targetModel,
+            geminiRequest,
+            token.token.access_token,
+          );
+          return this.processStreamResponse(stream, request.model);
+        } else {
+          const response = await this.geminiClient.generate(
+            targetModel,
+            geminiRequest,
+            token.token.access_token,
+          );
+          return this.convertGeminiToOpenAIResponse(response, request.model);
+        }
+      } catch (error) {
+        lastError = error;
+        if (error instanceof Error) {
+          this.logger.warn(`Request failed with account ${token.email}: ${error.message}`);
+
+          // Check for rate limit / auth errors
+          if (this.shouldRetry(error.message)) {
+            this.tokenManager.markAsRateLimited(token.email);
+          }
+        }
+      }
+    }
+    throw lastError || new Error('Request failed after retries');
+  }
+
+  // Handle SSE Stream conversion
+  private processStreamResponse(upstreamStream: any, model: string): Observable<string> {
+    return new Observable<string>((subscriber) => {
+      const decoder = new TextDecoder();
+      let buffer = '';
+
+      upstreamStream.on('data', (chunk: Buffer) => {
+        buffer += decoder.decode(chunk, { stream: true });
+        const lines = buffer.split('\n');
+        buffer = lines.pop() || '';
+
+        for (const line of lines) {
+          const trimmed = line.trim();
+          if (!trimmed.startsWith('data: ')) continue;
+
+          const dataStr = trimmed.slice(6);
+          if (dataStr === '[DONE]') continue;
+
+          try {
+            const json = JSON.parse(dataStr);
+            const candidate = json.candidates?.[0];
+            const contentPart = candidate?.content?.parts?.[0];
+            const text = contentPart?.text || '';
+
+            const openAICheck = {
+              id: `chatcmpl-${uuidv4()}`,
+              object: 'chat.completion.chunk',
+              created: Math.floor(Date.now() / 1000),
+              model: model,
+              choices: [
+                {
+                  index: 0,
+                  delta: { content: text },
+                  finish_reason: candidate?.finishReason ? 'stop' : null,
+                },
+              ],
+            };
+
+            subscriber.next(`data: ${JSON.stringify(openAICheck)}\n\n`);
+
+            if (candidate?.finishReason) {
+              subscriber.next('data: [DONE]\n\n');
+              subscriber.complete();
+            }
+          } catch (e) {
+            // ignore parse errors
+          }
+        }
+      });
+
+      upstreamStream.on('end', () => {
+        subscriber.complete();
+      });
+
+      upstreamStream.on('error', (err: any) => {
+        subscriber.error(err);
+      });
+    });
+  }
+
+  private convertOpenAIToGemini(request: OpenAIChatRequest): GeminiRequest {
+    const contents: GeminiContent[] = [];
+    let systemInstruction: { parts: GeminiPart[] } | undefined = undefined;
+
+    // Extract system message
+    const messages = request.messages || [];
+    for (const msg of messages) {
+      if (msg.role === 'system') {
+        if (typeof msg.content === 'string') {
+          systemInstruction = {
+            parts: [{ text: msg.content }],
+          };
+        }
+      } else {
+        let textPart = '';
+        if (Array.isArray(msg.content)) {
+          // OpenAI multimodal array
+          textPart = msg.content
+            .filter((p: OpenAIContentPart) => p.type === 'text')
+            .map((p: OpenAIContentPart) => p.text || '')
+            .join('\n');
+        } else {
+          textPart = msg.content;
+        }
+
+        contents.push({
+          role: msg.role === 'user' ? 'user' : 'model',
+          parts: [{ text: textPart }],
+        });
+      }
+    }
+
+    return {
+      contents,
+      systemInstruction,
+      generationConfig: {
+        temperature: request.temperature,
+        maxOutputTokens: request.max_tokens,
+        topP: request.top_p,
+      },
+    };
+  }
+
+  private convertGeminiToOpenAIResponse(
+    geminiResponse: GeminiResponse,
+    model: string,
+  ): OpenAIChatResponse {
+    const candidate = geminiResponse.candidates?.[0];
+    if (!candidate) throw new Error('No candidates in response');
+
+    // Process response for images
+    const processedCandidate = this.processInlineData(candidate);
+
+    const parts = processedCandidate?.content?.parts || [];
+    const content = parts.map((p) => p.text || '').join('');
+
+    const finishReason =
+      candidate?.finishReason === 'STOP'
+        ? 'stop'
+        : candidate?.finishReason?.toLowerCase() || 'stop';
+
+    return {
+      id: `chatcmpl-${uuidv4()}`,
+      object: 'chat.completion',
+      created: Math.floor(Date.now() / 1000),
+      model: model,
+      choices: [
+        {
+          index: 0,
+          message: {
+            role: 'assistant',
+            content: content,
+          },
+          finish_reason: finishReason,
+        },
+      ],
+      usage: {
+        prompt_tokens: geminiResponse.usageMetadata?.promptTokenCount || 0,
+        completion_tokens: geminiResponse.usageMetadata?.candidatesTokenCount || 0,
+        total_tokens: geminiResponse.usageMetadata?.totalTokenCount || 0,
+      },
+    };
+  }
+
+  // Convert Gemini inlineData (images) to Markdown
+  private processInlineData(candidate: GeminiCandidate | undefined): GeminiCandidate | undefined {
+    if (!candidate?.content?.parts) return candidate;
+
+    const newParts = candidate.content.parts.map((part) => {
+      if (part.inlineData) {
+        const mimeType = part.inlineData.mimeType || 'image/jpeg';
+        const data = part.inlineData.data || '';
+        return {
+          text: `\n\n![Generated Image](data:${mimeType};base64,${data})\n\n`,
+        };
+      }
+      return part;
+    });
+
+    return {
+      ...candidate,
+      content: {
+        ...candidate.content,
+        parts: newParts,
+      },
+    };
+  }
+
+  private mapModel(model: string): string {
+    const lower = model.toLowerCase();
+    if (lower.includes('sonnet') || lower.includes('thinking')) return 'gemini-3-pro-preview';
+    if (lower.includes('haiku')) return 'gemini-2.0-flash-exp';
+    if (lower.includes('opus')) return 'gemini-3-pro-preview';
+    if (lower.includes('claude')) return 'gemini-2.5-flash-thinking';
+    if (lower === 'gemini-3-pro-high' || lower === 'gemini-3-pro-low')
+      return 'gemini-3-pro-preview';
+    if (lower === 'gemini-3-flash') return 'gemini-3-flash-preview';
+
+    // Default to original if no mapping found
+    return model;
+  }
+
+  private shouldRetry(errorMessage: string): boolean {
+    const msg = errorMessage.toLowerCase();
+    if (
+      msg.includes('404') ||
+      msg.includes('not_found') ||
+      msg.includes('403') ||
+      msg.includes('permission_denied')
+    )
+      return true;
+    if (
+      msg.includes('429') ||
+      msg.includes('resource_exhausted') ||
+      msg.includes('quota') ||
+      msg.includes('rate_limit')
+    )
+      return true;
+    return false;
+  }
+}

src/server/modules/proxy/token-manager.service.ts

@@ -0,0 +1,202 @@
+import { Injectable, Logger, OnModuleInit } from '@nestjs/common';
+import { CloudAccountRepo } from '../../../ipc/database/cloudHandler';
+import { CloudAccount } from '../../../types/cloudAccount';
+import { GoogleAPIService } from '../../../services/GoogleAPIService';
+
+interface TokenData {
+  email: string;
+  account_id: string;
+  access_token: string;
+  refresh_token: string;
+  expires_in: number;
+  expiry_timestamp: number;
+  project_id?: string;
+  session_id?: string;
+}
+
+@Injectable()
+export class TokenManagerService implements OnModuleInit {
+  private readonly logger = new Logger(TokenManagerService.name);
+  private currentIndex = 0;
+  // In-memory cache of tokens with additional data
+  private tokens: Map<string, TokenData> = new Map();
+  // Cooldown map for rate-limited accounts
+  private cooldowns: Map<string, number> = new Map();
+
+  async onModuleInit() {
+    // Load accounts on module initialization
+    await this.loadAccounts();
+  }
+
+  async loadAccounts(): Promise<number> {
+    try {
+      const accounts = await CloudAccountRepo.getAccounts();
+      let count = 0;
+
+      for (const account of accounts) {
+        const tokenData = this.convertAccountToToken(account);
+        if (tokenData) {
+          this.tokens.set(account.id, tokenData);
+          count++;
+        }
+      }
+
+      this.logger.log(`Loaded ${count} accounts`);
+      return count;
+    } catch (e) {
+      this.logger.error('Failed to load accounts', e);
+      return 0;
+    }
+  }
+
+  private convertAccountToToken(account: CloudAccount): TokenData | null {
+    if (!account.token) return null;
+
+    return {
+      account_id: account.id,
+      email: account.email,
+      access_token: account.token.access_token,
+      refresh_token: account.token.refresh_token,
+      expires_in: account.token.expires_in,
+      expiry_timestamp: account.token.expiry_timestamp,
+      project_id: account.token.project_id || undefined,
+      session_id: account.token.session_id || this.generateSessionId(),
+    };
+  }
+
+  private generateSessionId(): string {
+    // Match the Rust implementation's session ID format
+    const min = 1_000_000_000_000_000_000n;
+    const max = 9_000_000_000_000_000_000n;
+    const range = max - min;
+    const rand = BigInt(Math.floor(Math.random() * Number(range)));
+    return (-(min + rand)).toString();
+  }
+
+  async getNextToken(): Promise<CloudAccount | null> {
+    try {
+      // Reload if empty
+      if (this.tokens.size === 0) {
+        await this.loadAccounts();
+      }
+      if (this.tokens.size === 0) return null;
+
+      const now = Date.now();
+      const nowSeconds = Math.floor(now / 1000);
+
+      // Filter out accounts in cooldown
+      const validTokens = Array.from(this.tokens.entries()).filter(([email, _]) => {
+        const cooldownUntil = this.cooldowns.get(email);
+        return !cooldownUntil || cooldownUntil <= now;
+      });
+
+      if (validTokens.length === 0) {
+        this.logger.warn('All accounts are in cooldown');
+        return null;
+      }
+
+      // Round robin selection
+      const [accountId, tokenData] = validTokens[this.currentIndex % validTokens.length];
+      this.currentIndex++;
+
+      // Check if token needs refresh (expires in < 5 minutes)
+      if (nowSeconds >= tokenData.expiry_timestamp - 300) {
+        this.logger.log(`Token for ${tokenData.email} expiring soon, refreshing...`);
+        try {
+          const newTokens = await GoogleAPIService.refreshAccessToken(tokenData.refresh_token);
+
+          // Update token data
+          tokenData.access_token = newTokens.access_token;
+          tokenData.expires_in = newTokens.expires_in;
+          tokenData.expiry_timestamp = nowSeconds + newTokens.expires_in;
+
+          // Save to DB
+          await this.saveRefreshedToken(accountId, tokenData);
+          this.tokens.set(accountId, tokenData);
+
+          this.logger.log(`Token refreshed for ${tokenData.email}`);
+        } catch (e) {
+          this.logger.error(`Failed to refresh token for ${tokenData.email}`, e);
+        }
+      }
+
+      // Resolve project ID if missing (mock for now, like original)
+      if (!tokenData.project_id) {
+        const mockId = `cloud-code-${Math.floor(Math.random() * 100000)}`;
+        tokenData.project_id = mockId;
+        await this.saveProjectId(accountId, mockId);
+      }
+
+      this.logger.log(`Selected account: ${tokenData.email}`);
+
+      // Return in CloudAccount format for compatibility
+      return {
+        id: accountId,
+        email: tokenData.email,
+        token: {
+          access_token: tokenData.access_token,
+          refresh_token: tokenData.refresh_token,
+          expires_in: tokenData.expires_in,
+          expiry_timestamp: tokenData.expiry_timestamp,
+          project_id: tokenData.project_id,
+          session_id: tokenData.session_id,
+        },
+      } as CloudAccount;
+    } catch (error) {
+      this.logger.error('Failed to get token', error);
+      return null;
+    }
+  }
+
+  markAsRateLimited(email: string) {
+    // Cooldown for 5 minutes
+    const until = Date.now() + 5 * 60 * 1000;
+    this.cooldowns.set(email, until);
+    this.logger.warn(
+      `Account ${email} marked as rate limited until ${new Date(until).toISOString()}`,
+    );
+  }
+
+  resetCooldown(email: string) {
+    this.cooldowns.delete(email);
+  }
+
+  private async saveRefreshedToken(accountId: string, tokenData: TokenData) {
+    try {
+      const acc = await CloudAccountRepo.getAccount(accountId);
+      if (acc && acc.token) {
+        const newToken = {
+          ...acc.token,
+          access_token: tokenData.access_token,
+          expires_in: tokenData.expires_in,
+          expiry_timestamp: tokenData.expiry_timestamp,
+        };
+        await CloudAccountRepo.updateToken(accountId, newToken);
+      }
+    } catch (e) {
+      this.logger.error('Failed to save refreshed token to DB', e);
+    }
+  }
+
+  private async saveProjectId(accountId: string, projectId: string) {
+    try {
+      const acc = await CloudAccountRepo.getAccount(accountId);
+      if (acc && acc.token) {
+        const newToken = {
+          ...acc.token,
+          project_id: projectId,
+        };
+        await CloudAccountRepo.updateToken(accountId, newToken);
+      }
+    } catch (e) {
+      this.logger.error('Failed to save project ID to DB', e);
+    }
+  }
+
+  /**
+   * Get the number of loaded accounts (for status)
+   */
+  getAccountCount(): number {
+    return this.tokens.size;
+  }
+}