main.ts

import { arch, env, platform } from "node:process";
import {
  AuthType,
  createCodeAssistContentGenerator,
} from "npm:@google/gemini-cli-core@0.1.4";
import { GaxiosError } from "npm:gaxios@6.7.1";

env.CLI_VERSION ||= "0.1.4";

const codeAssist = await createCodeAssistContentGenerator(
  {
    headers: {
      "User-Agent": `GeminiCLI/${env.CLI_VERSION} (${platform}; ${arch})`,
    },
  },
  AuthType.LOGIN_WITH_GOOGLE_PERSONAL,
);

Deno.serve(async (req) => {
  const pathname = new URL(req.url).pathname;
  const [model, action] = pathname
    .split("/")
    .find((part) => part.includes(":"))
    ?.split(":") ?? [];

  if (!model || !action) {
    return new Response("Invalid request", {
      status: 400,
    });
  }

  const payload = await req.json().catch(() => ({}));

  // Uncomment for debugging
  console.info(model, action, payload);

  const getGenerateContentParameters = () => ({
    model,
    contents: payload.contents,
    config: {
      ...payload.generationConfig,
      tools: payload.tools,
      toolConfig: payload.toolConfig,
      safetySettings: payload.safetySettings,
      systemInstruction: payload.systemInstruction,
      abortSignal: req.signal,
    },
  } as const);

  try {
    switch (action) {
      case "generateContent": {
        const result = await codeAssist.generateContent(
          getGenerateContentParameters(),
        );
        return Response.json(result);
      }
      case "streamGenerateContent": {
        const stream = await codeAssist.generateContentStream(
          getGenerateContentParameters(),
        );
        
        // Create a ReadableStream that processes the Gemini stream
        const readableStream = new ReadableStream({
          async start(controller) {
            try {
              for await (const chunk of stream) {
                try {
                  const jsonData = JSON.stringify(chunk);
                  // Send SSE formatted data
                  const sseData = `data: ${jsonData}\n\n`;
                  controller.enqueue(new TextEncoder().encode(sseData));
                } catch (jsonError) {
                  console.error("Error serializing chunk:", jsonError, chunk);
                  const errorData = `data: ${JSON.stringify({ error: "Serialization error" })}\n\n`;
                  controller.enqueue(new TextEncoder().encode(errorData));
                }
              }
            } catch (error) {
              console.error("Error in stream:", error);
              const errorData = `data: ${JSON.stringify({ error: "Stream processing error" })}\n\n`;
              controller.enqueue(new TextEncoder().encode(errorData));
            } finally {
              console.log("Closing stream");
              controller.close();
            }
          }
        });

        return new Response(readableStream, {
          headers: {
            "Content-Type": "text/event-stream",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Content-Type",
          },
        });
      }
      case "countTokens": {
        const result = await codeAssist.countTokens({
          model,
          contents: payload.contents,
          config: {
            ...payload.generateContentRequest,
            abortSignal: req.signal,
          },
        });
        return Response.json(result);
      }
      case "embedContent": {
        const result = await codeAssist.embedContent({
          model,
          contents: payload.contents,
          config: {
            taskType: payload.taskType,
            title: payload.title,
            outputDimensionality: payload.outputDimensionality,
            abortSignal: req.signal,
          },
        });
        return Response.json(result);
      }
      default: {
        return new Response(`Invalid action: ${action}`, {
          status: 400,
        });
      }
    }
  } catch (error) {
    if (error instanceof GaxiosError && error.response) {
      return Response.json(error.response.data, {
        status: error.response.status,
        statusText: error.response.statusText,
        headers: error.response.headers,
      });
    }
    console.error("Error processing request:", error);
    return new Response("Internal Server Error", {
      status: 500,
    });
  }
});