From eaa35168ea74939f18620b7ab4fb342915ecb1b6 Mon Sep 17 00:00:00 2001
From: Shawn Morreau <shawn@arktype.io>
Date: Wed, 26 Nov 2025 15:14:18 -0500
Subject: [PATCH] gemini retries

---
 agents/gemini.ts |  4 ++--
 agents/shared.ts | 58 ++++++++++++++++++++++++++++++++++--------------
 main.ts          |  7 +++++-
 3 files changed, 49 insertions(+), 20 deletions(-)

diff --git a/agents/gemini.ts b/agents/gemini.ts
index cc3eb27..4ca7f10 100644
--- a/agents/gemini.ts
+++ b/agents/gemini.ts
@@ -144,12 +144,12 @@ const messageHandlers = {
 
 export const gemini = agent({
   name: "gemini",
-  install: async () => {
+  install: async (githubInstallationToken?: string) => {
     return await installFromGithub({
       owner: "google-gemini",
       repo: "gemini-cli",
-      tag: "v0.16.0",
       assetName: "gemini.js",
+      ...(githubInstallationToken && { githubInstallationToken }),
     });
   },
   run: async ({ payload, apiKey, mcpServers, githubInstallationToken, cliPath }) => {
diff --git a/agents/shared.ts b/agents/shared.ts
index 24620d8..1020fc4 100644
--- a/agents/shared.ts
+++ b/agents/shared.ts
@@ -62,9 +62,9 @@ export interface InstallFromCurlParams {
 export interface InstallFromGithubParams {
   owner: string;
   repo: string;
-  tag?: string;
   assetName?: string;
   executablePath?: string;
+  githubInstallationToken?: string;
 }
 
 /**
@@ -180,6 +180,36 @@ export async function installFromNpmTarball({
   return cliPath;
 }
 
+/**
+ * Fetch with retry logic if Retry-After header is present
+ */
+async function fetchWithRetry(
+  url: string,
+  headers: Record<string, string>,
+  errorMessage: string
+): Promise<Response> {
+  const response = await fetch(url, { headers });
+  if (!response.ok) {
+    const retryAfter = response.headers.get("Retry-After") || response.headers.get("retry-after");
+    if (retryAfter) {
+      const waitSeconds = parseInt(retryAfter, 10);
+      if (!Number.isNaN(waitSeconds) && waitSeconds > 0) {
+        log.info(`Rate limited, waiting ${waitSeconds} seconds before retry...`);
+        await new Promise((resolve) => setTimeout(resolve, waitSeconds * 1000));
+        const retryResponse = await fetch(url, { headers });
+        if (!retryResponse.ok) {
+          throw new Error(
+            `${errorMessage}: ${retryResponse.status} ${retryResponse.statusText} (retry failed)`
+          );
+        }
+        return retryResponse;
+      }
+    }
+    throw new Error(`${errorMessage}: ${response.status} ${response.statusText}`);
+  }
+  return response;
+}
+
 /**
  * Install a CLI tool from GitHub releases
  * Downloads the latest release asset from GitHub and returns the path to the executable
@@ -188,24 +218,23 @@ export async function installFromNpmTarball({
 export async function installFromGithub({
   owner,
   repo,
-  tag,
   assetName,
   executablePath,
+  githubInstallationToken,
 }: InstallFromGithubParams): Promise<string> {
   log.info(`📦 Installing ${owner}/${repo} from GitHub releases...`);
 
-  // fetch release from GitHub API (specific tag or latest)
-  const releaseUrl = tag
-    ? `https://api.github.com/repos/${owner}/${repo}/releases/tags/${tag}`
-    : `https://api.github.com/repos/${owner}/${repo}/releases/latest`;
+  // fetch release from GitHub API (latest)
+  const releaseUrl = `https://api.github.com/repos/${owner}/${repo}/releases/latest`;
   log.info(`Fetching release from ${releaseUrl}...`);
-  const releaseResponse = await fetch(releaseUrl);
-  if (!releaseResponse.ok) {
-    throw new Error(
-      `Failed to fetch release: ${releaseResponse.status} ${releaseResponse.statusText}`
-    );
+
+  const headers: Record<string, string> = {};
+  if (githubInstallationToken) {
+    headers.Authorization = `Bearer ${githubInstallationToken}`;
   }
 
+  const releaseResponse = await fetchWithRetry(releaseUrl, headers, "Failed to fetch release");
+
   const releaseData = (await releaseResponse.json()) as {
     tag_name: string;
     assets: Array<{
@@ -234,12 +263,7 @@ export async function installFromGithub({
   const downloadPath = join(tempDir, fileName);
 
   // download the asset
-  const assetResponse = await fetch(assetUrl);
-  if (!assetResponse.ok) {
-    throw new Error(
-      `Failed to download asset: ${assetResponse.status} ${assetResponse.statusText}`
-    );
-  }
+  const assetResponse = await fetchWithRetry(assetUrl, headers, "Failed to download asset");
 
   if (!assetResponse.body) throw new Error("Response body is null");
   const fileStream = createWriteStream(downloadPath);
diff --git a/main.ts b/main.ts
index b3e3c87..e6f9ad8 100644
--- a/main.ts
+++ b/main.ts
@@ -300,7 +300,12 @@ function setupMcpServers(ctx: MainContext): void {
 }
 
 async function installAgentCli(ctx: MainContext): Promise<void> {
-  ctx.cliPath = await ctx.agent.install();
+  // gemini is the only agent that needs githubInstallationToken for install
+  if (ctx.agentName === "gemini") {
+    ctx.cliPath = await ctx.agent.install(ctx.githubInstallationToken);
+  } else {
+    ctx.cliPath = await ctx.agent.install();
+  }
 }
 
 function validateApiKey(ctx: MainContext): void {