Fix image retry loop

[starkdmi]

Fix for 164, resolves account rotation on quota limits reached with image model.

[coderabbitai]

Walkthrough

The changes implement header-style-aware rate limiting and expand explicit quota handling for image models. The getMinWaitTimeForFamily method signature is updated to accept headerStyle and a strict boolean parameter for finer-grained rate-limit checking. The accounts module now conditionally uses header-style-specific rate-limit queries when strict mode is enabled. The model resolver expands explicit quota marking to include image generation models alongside Antigravity models. Tests are added to verify header-style-specific rate-limit behavior and image model classification.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Fix image retry loop' is directly related to the main purpose of the PR, which addresses account rotation issues when quota limits are reached for image models. The changeset focuses on fixing how wait times are calculated for image model quotas to enable proper account rotation.
Description check	✅ Passed	The description clearly references issue/PR 164 and explains the fix resolves account rotation on quota limits for image models, which directly relates to the changeset modifications involving image model quota handling and wait time calculations.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

📜 Recent review details

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to data retention organization setting

📥 Commits

Reviewing files that changed from the base of the PR and between 7f73cb2 and 973dfbb.

📒 Files selected for processing (5)

• src/plugin.ts
• src/plugin/accounts.test.ts
• src/plugin/accounts.ts
• src/plugin/transform/model-resolver.test.ts
• src/plugin/transform/model-resolver.ts

🧰 Additional context used

🧬 Code graph analysis (3)

src/plugin/accounts.test.ts (1)

src/plugin/storage.ts (1)

• AccountStorageV3 (190-198)

src/plugin/transform/model-resolver.test.ts (2)

src/plugin/transform/model-resolver.ts (1)

• resolveModelWithTier (177-282)

src/plugin/transform/index.ts (1)

• resolveModelWithTier (22-22)

src/plugin/accounts.ts (3)

src/plugin/transform/types.ts (1)

• ModelFamily (3-3)

src/plugin/storage.ts (2)

• ModelFamily (125-125)
• HeaderStyle (126-126)

src/constants.ts (1)

• HeaderStyle (85-85)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Greptile Review

🔇 Additional comments (5)

src/plugin/transform/model-resolver.test.ts (1)

76-92: LGTM - Good test coverage for image model quota handling.

The tests properly validate that both antigravity-gemini-3-pro-image and gemini-3-pro-image resolve with explicitQuota: true and quotaPreference: "antigravity", which aligns with the fix for the image retry loop issue.

src/plugin/transform/model-resolver.ts (1)

188-190: LGTM - Correct fix for image model quota handling.

Setting explicitQuota = isAntigravity || isImageModel ensures that image models participate in the strict header-style-aware wait time calculations. Since image models always route to Antigravity (line 189), marking them as explicit quota is the correct behavior to prevent the retry loop.

src/plugin/accounts.test.ts (1)

1141-1168: LGTM - Test correctly validates strict header-style wait time calculation.

This test effectively demonstrates the core fix: when strict mode is enabled with a specific headerStyle, getMinWaitTimeForFamily returns the wait time for that specific quota (30000ms) rather than falling back to 0 because another quota pool might be available. This is essential for preventing the image model retry loop.

src/plugin.ts (1)

906-914: LGTM - Correctly passes headerStyle and explicitQuota for stricter wait time calculation.

The fix appropriately uses explicitQuota (derived from the model) as the strict parameter. For image models (where explicitQuota=true), this ensures the wait time is calculated strictly for the Antigravity quota, preventing the retry loop that occurred when the system incorrectly assumed another quota pool was available.

src/plugin/accounts.ts (1)

615-657: LGTM - Well-implemented strict mode for header-style-aware wait time calculation.

The implementation correctly differentiates between:

1. Strict mode (lines 623-625, 636-639): When strict && headerStyle, checks and calculates wait times for the specific quota key only. This is essential for image models that must use the Antigravity quota.
2. Non-strict mode (lines 640-652): Preserves existing behavior where Gemini family accounts become available when either quota pool expires.

This separation is the core fix for the image retry loop - image models with explicitQuota=true now correctly wait for their specific quota to reset rather than incorrectly reporting 0ms wait time because a different quota pool might be available.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[greptile-apps]

Greptile Summary

Fixes infinite retry loop for image generation models when quota limits are reached. Previously, when image models hit the antigravity quota limit, getMinWaitTimeForFamily would incorrectly return 0 by falling back to the generic Gemini quota logic that checks EITHER pool (antigravity OR gemini-cli). Since image models exclusively use antigravity, this caused immediate retries instead of waiting for quota reset.

The fix adds strict quota tracking:

• Marks image models with explicitQuota=true in model-resolver.ts
• Passes headerStyle and explicitQuota parameters to getMinWaitTimeForFamily in plugin.ts
• Adds strict mode to getMinWaitTimeForFamily that respects the specific quota pool when strict=true
• Includes comprehensive test coverage validating the strict behavior

Confidence Score: 5/5

• Safe to merge with minimal risk
• The fix is targeted, well-tested, and directly addresses the root cause. Changes are minimal and isolated to the specific issue - adding strict quota tracking for image models without affecting other model types. Test coverage validates both strict and non-strict modes.
• No files require special attention

Important Files Changed

Filename	Overview
src/plugin/transform/model-resolver.ts	Marks image models as explicitQuota=true to enable strict quota tracking
src/plugin.ts	Passes headerStyle and explicitQuota to getMinWaitTimeForFamily for strict wait time calculation
src/plugin/accounts.ts	Adds strict mode to getMinWaitTimeForFamily to respect specific quota pool when explicitQuota=true

Sequence Diagram

sequenceDiagram
    participant Client
    participant Plugin as plugin.ts
    participant Resolver as model-resolver.ts
    participant AccountMgr as AccountManager

    Note over Client,AccountMgr: Image Model Request (e.g., gemini-3-pro-image)
    
    Client->>Plugin: Request with image model
    Plugin->>Resolver: resolveModelWithTier(model)
    Resolver->>Resolver: Detect IMAGE_GENERATION_MODELS regex
    Resolver->>Resolver: Set explicitQuota=true (NEW)
    Resolver-->>Plugin: {explicitQuota: true, quotaPreference: "antigravity"}
    
    Plugin->>AccountMgr: getCurrentOrNextForFamily()
    AccountMgr->>AccountMgr: Check rate limits
    
    alt Account available
        AccountMgr-->>Plugin: Return account
        Plugin->>Client: Process request
    else All accounts rate-limited
        AccountMgr-->>Plugin: null
        Plugin->>Plugin: getHeaderStyleFromUrl() → "antigravity"
        Plugin->>Plugin: isExplicitQuotaFromUrl() → true
        Plugin->>AccountMgr: getMinWaitTimeForFamily(family, model, "antigravity", true)
        Note over AccountMgr: NEW: strict=true, headerStyle="antigravity"
        AccountMgr->>AccountMgr: Filter by isRateLimitedForHeaderStyle<br/>(strict quota check)
        AccountMgr->>AccountMgr: Calculate wait time for<br/>specific antigravity quota pool
        AccountMgr-->>Plugin: waitMs (e.g., 30000)
        Plugin->>Client: Return 429 with retry-after
    end
    
    Note over Plugin,AccountMgr: Before fix: Used fallback logic (EITHER pool)<br/>After fix: Respects strict antigravity quota

Loading

[greptile-apps]

Greptile found no issues!

From now on, if a review finishes and we haven't found any issues, we will not post anything, but you can confirm that we reviewed your changes in the status check section.

_{This feature can be toggled off in your Code Review Settings by deselecting "Create a status check for each PR".}