feat: Add canUpgradePlan field to quota-related classes to guard upgrade plan button. (#212)

Author: ethanyhou

com.microsoft.copilot.eclipse.core/src/com/microsoft/copilot/eclipse/core/lsp/protocol/quota/CheckQuotaResult.java

@@ -13,6 +13,8 @@
  * @param resetDateUtc ISO-8601 instant when the monthly allowance resets in UTC, or {@code null}
  * @param copilotPlan the user's Copilot plan
  * @param tokenBasedBillingEnabled whether the user's billing is token-based
+ * @param canUpgradePlan whether the user is eligible to upgrade their Copilot plan; {@code null} when the language
+ *     server has not supplied this field, in which case callers should fall back to plan-based defaults
  */
 public record CheckQuotaResult(
     Quota chat,
@@ -21,10 +23,11 @@ public record CheckQuotaResult(
     String resetDate,
     String resetDateUtc,
     CopilotPlan copilotPlan,
-    boolean tokenBasedBillingEnabled) {
+    boolean tokenBasedBillingEnabled,
+    Boolean canUpgradePlan) {
 
   private static final CheckQuotaResult EMPTY =
-      new CheckQuotaResult(null, null, null, null, null, null, false);
+      new CheckQuotaResult(null, null, null, null, null, null, false, null);
 
   /**
    * Returns an empty {@link CheckQuotaResult} used as a placeholder before the language server

com.microsoft.copilot.eclipse.core/src/com/microsoft/copilot/eclipse/core/lsp/protocol/quota/QuotaWarningParams.java

@@ -16,8 +16,11 @@
  * @param completions current completions quota snapshot, when available
  * @param premiumInteractions current premium interactions quota snapshot, when available
  * @param copilotPlan the user's Copilot plan
+ * @param canUpgradePlan whether the user is eligible to upgrade their Copilot plan; {@code null} when the language
+ *     server has not supplied this field, in which case callers should fall back to plan-based defaults
  */
 public record QuotaWarningParams(String title, String message, String severity, QuotaSnapshotParams chat,
     QuotaSnapshotParams completions,
-    @SerializedName("premium_interactions") QuotaSnapshotParams premiumInteractions, CopilotPlan copilotPlan) {
+    @SerializedName("premium_interactions") QuotaSnapshotParams premiumInteractions, CopilotPlan copilotPlan,
+    Boolean canUpgradePlan) {
 }

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/chat/ActionBar.java

@@ -977,16 +977,20 @@ public void createRateLimitBanner(String message, boolean warning) {
 
   /**
    * Show the quota-warning static banner above the input area. Action links are sourced from
-   * {@link QuotaActions#forPlan(CopilotPlan, boolean)} so they stay in sync with the inline {@link WarnWidget}.
+   * {@link QuotaActions#forPlan(CopilotPlan, boolean, Boolean)} so they stay in sync with the inline
+   * {@link WarnWidget}.
    *
    * @param message the message to display
    * @param plan the user's Copilot plan, or {@code null} for no action links
    * @param overageEnabled whether additional paid usage is already enabled for the user; switches the
    *     "Enable Additional Usage" label to "Increase Budget"
+   * @param canUpgradePlan whether the user can upgrade their Copilot plan, or {@code null} when the language
+   *     server did not supply this field
    * @param warning {@code true} for the warning icon; {@code false} for the info icon
    */
-  public void createQuotaWarningBanner(String message, CopilotPlan plan, boolean overageEnabled, boolean warning) {
-    List<BannerAction> bannerActions = QuotaActions.forPlan(plan, overageEnabled).stream()
+  public void createQuotaWarningBanner(String message, CopilotPlan plan, boolean overageEnabled,
+      Boolean canUpgradePlan, boolean warning) {
+    List<BannerAction> bannerActions = QuotaActions.forPlan(plan, overageEnabled, canUpgradePlan).stream()
         .map(action -> new BannerAction(action.label(), action.url()))
         .toList();
     showStaticBanner(message, bannerActions, warning);

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/chat/BaseTurnWidget.java

@@ -28,6 +28,7 @@
 import com.microsoft.copilot.eclipse.core.lsp.protocol.AgentToolCall;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.LanguageModelToolConfirmationResult;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.codingagent.CodingAgentMessageRequestParams;
+import com.microsoft.copilot.eclipse.core.lsp.protocol.quota.CheckQuotaResult;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.quota.CopilotPlan;
 import com.microsoft.copilot.eclipse.core.persistence.ConversationDataFactory;
 import com.microsoft.copilot.eclipse.core.persistence.CopilotTurnData;
@@ -611,13 +612,15 @@ protected void createWarnDialog(String message, int code, String modelProviderNa
     String displayMessage = byokQuotaExceeded ? Messages.chat_warnWidget_byokQuotaUsageMessage : message;
     CopilotPlan planForActions = null;
     boolean overageEnabled = false;
+    Boolean canUpgradePlan = null;
     if (code == 402 && !byokQuotaExceeded) {
-      var quotaStatus = this.serviceManager.getAuthStatusManager().getQuotaStatus();
+      CheckQuotaResult quotaStatus = this.serviceManager.getAuthStatusManager().getQuotaStatus();
       planForActions = quotaStatus.copilotPlan();
       overageEnabled = quotaStatus.premiumInteractions() != null
           && quotaStatus.premiumInteractions().overagePermitted();
+      canUpgradePlan = quotaStatus.canUpgradePlan();
     }
-    new WarnWidget(this, SWT.NONE, displayMessage, planForActions, overageEnabled);
+    new WarnWidget(this, SWT.NONE, displayMessage, planForActions, overageEnabled, canUpgradePlan);
     ensureFooterAtBottom();
     requestLayout();
   }

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/chat/ChatView.java

@@ -386,7 +386,8 @@ public void done(IJobChangeEvent event) {
             boolean warning = "warning".equalsIgnoreCase(params.severity());
             boolean overageEnabled = params.premiumInteractions() != null
                 && params.premiumInteractions().overageEnabled();
-            actionBar.createQuotaWarningBanner(params.message(), params.copilotPlan(), overageEnabled, warning);
+            actionBar.createQuotaWarningBanner(params.message(), params.copilotPlan(), overageEnabled,
+                params.canUpgradePlan(), warning);
           }
         }, parent);
       }

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/chat/QuotaActions.java

@@ -37,24 +37,31 @@ private QuotaActions() {
   /**
    * Returns the ordered list of {@link QuotaAction}s appropriate for the supplied plan.
    *
-   * <p>Mapping:
+   * <p>Non-upgrade actions are derived from the plan:
    * <ul>
-   *   <li>{@code free} &rarr; "Upgrade Plan" (primary)</li>
-   *   <li>{@code individual}, {@code individual_pro} &rarr; "Enable Additional Usage" /
-   *       "Increase Budget" (primary) + "Upgrade Plan" (secondary)</li>
-   *   <li>{@code individual_max} &rarr; "Enable Additional Usage" / "Increase Budget" (primary)</li>
-   *   <li>{@code business}, {@code enterprise} &rarr; empty list</li>
-   *   <li>{@code null} &rarr; empty list</li>
+   *   <li>{@code free} &rarr; none</li>
+   *   <li>{@code individual}, {@code individual_pro}, {@code individual_max} &rarr;
+   *       "Enable Additional Usage" / "Increase Budget"</li>
+   *   <li>{@code business}, {@code enterprise} &rarr; none</li>
    * </ul>
    *
-   * <p>The "Enable Additional Usage" label is replaced with "Increase Budget" when {@code overageEnabled}
-   * is {@code true}, matching the IntelliJ quota dialog wording.
+   * <p>The "Upgrade Plan" action is then appended when the user is eligible. When
+   * {@code canUpgradePlan} is non-{@code null} it takes precedence over the plan-based default;
+   * otherwise the plan default is used (eligible for {@code free}, {@code individual},
+   * {@code individual_pro}). When the upgrade action is the only entry it is rendered as primary;
+   * when it follows another primary action it is rendered as secondary so the surfaces stay
+   * visually consistent.
+   *
+   * <p>The "Enable Additional Usage" label is replaced with "Increase Budget" when
+   * {@code overageEnabled} is {@code true}, matching the IntelliJ quota dialog wording.
    *
    * @param plan the user's Copilot plan, or {@code null} when unknown
    * @param overageEnabled {@code true} when additional paid usage is already enabled for the user
+   * @param canUpgradePlan whether the user can upgrade their Copilot plan, or {@code null} when the
+   *     language server did not supply this field
    * @return an immutable, possibly empty list; never {@code null}
    */
-  public static List<QuotaAction> forPlan(CopilotPlan plan, boolean overageEnabled) {
+  public static List<QuotaAction> forPlan(CopilotPlan plan, boolean overageEnabled, Boolean canUpgradePlan) {
     if (plan == null) {
       return List.of();
     }
@@ -69,20 +76,34 @@ public static List<QuotaAction> forPlan(CopilotPlan plan, boolean overageEnabled
     QuotaAction manageOverage = new QuotaAction(overageLabel,
         Messages.chat_noQuotaView_enableAdditionalUsageButton_tooltip,
         UiConstants.MANAGE_COPILOT_OVERAGE_URL, true);
-    return switch (plan) {
-      case free -> List.of(upgradePrimary);
-      case individual, individual_pro -> List.of(manageOverage, upgradeSecondary);
-      case individual_max -> List.of(manageOverage);
-      case business, enterprise -> List.of();
+
+    boolean hasOverage = switch (plan) {
+      case individual, individual_pro, individual_max -> true;
+      case free, business, enterprise -> false;
+    };
+    boolean showUpgrade = canUpgradePlan != null ? canUpgradePlan : switch (plan) {
+      case free, individual, individual_pro -> true;
+      case individual_max, business, enterprise -> false;
     };
+    if (hasOverage && showUpgrade) {
+      return List.of(manageOverage, upgradeSecondary);
+    }
+    if (hasOverage) {
+      return List.of(manageOverage);
+    }
+    if (showUpgrade) {
+      return List.of(upgradePrimary);
+    }
+    return List.of();
   }
 
   /**
    * Returns {@code true} when an error response represents a Bring-Your-Own-Key (BYOK) quota-exceeded
    * condition, i.e. a {@code 402} from the language server that also carries a non-blank model
    * provider name. BYOK usage is governed by the customer's provider account rather than the user's
-   * Copilot plan, so callers should suppress the plan-driven {@link #forPlan(CopilotPlan)} actions
-   * and substitute the BYOK-specific message in this case.
+   * Copilot plan, so callers should suppress the plan-driven
+   * {@link #forPlan(CopilotPlan, boolean, Boolean)} actions and substitute the BYOK-specific
+   * message in this case.
    *
    * @param code the error code from the language server
    * @param modelProviderName the BYOK model-provider name from the server payload, or {@code null}

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/chat/WarnWidget.java

@@ -26,8 +26,8 @@
 
 /**
  * Widget that displays a warning message under a chat turn, optionally followed by plan-driven action buttons sourced
- * from {@link QuotaActions#forPlan(CopilotPlan, boolean)}. Presentation-only: the caller decides the message and
- * whether to pass a plan.
+ * from {@link QuotaActions#forPlan(CopilotPlan, boolean, Boolean)}. Presentation-only: the caller decides the message
+ * and whether to pass a plan.
  */
 public class WarnWidget extends Composite {
   private int buttonLeftMargin;
@@ -41,8 +41,11 @@ public class WarnWidget extends Composite {
    * @param userPlan the user's Copilot plan to render plan-driven action buttons, or {@code null} for no buttons
    * @param overageEnabled whether additional paid usage is already enabled for the user; switches the
    *     "Enable Additional Usage" label to "Increase Budget"
+   * @param canUpgradePlan whether the user can upgrade their Copilot plan, or {@code null} when the language
+   *     server did not supply this field; forwarded to {@link QuotaActions#forPlan(CopilotPlan, boolean, Boolean)}
    */
-  public WarnWidget(Composite parent, int style, String message, CopilotPlan userPlan, boolean overageEnabled) {
+  public WarnWidget(Composite parent, int style, String message, CopilotPlan userPlan, boolean overageEnabled,
+      Boolean canUpgradePlan) {
     super(parent, style | SWT.BORDER);
     GridLayout outerLayout = new GridLayout(1, true);
     outerLayout.verticalSpacing = 0;
@@ -52,7 +55,7 @@ public WarnWidget(Composite parent, int style, String message, CopilotPlan userP
     buildWarnLabelWithIcon(message);
 
     if (userPlan != null) {
-      buildActionButtons(userPlan, overageEnabled);
+      buildActionButtons(userPlan, overageEnabled, canUpgradePlan);
     }
     parent.layout();
   }
@@ -128,8 +131,8 @@ private void buildWarnLabelWithIcon(String message) {
   /**
    * Render plan-driven action buttons for a quota-exceeded warning, kept in sync with the quota {@link StaticBanner}.
    */
-  private void buildActionButtons(CopilotPlan userPlan, boolean overageEnabled) {
-    List<QuotaAction> actions = QuotaActions.forPlan(userPlan, overageEnabled);
+  private void buildActionButtons(CopilotPlan userPlan, boolean overageEnabled, Boolean canUpgradePlan) {
+    List<QuotaAction> actions = QuotaActions.forPlan(userPlan, overageEnabled, canUpgradePlan);
     if (actions.isEmpty()) {
       return;
     }

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/handlers/ShowMenuBarMenuHandler.java

@@ -234,7 +234,7 @@ private void addCopilotUsageItemsTbb(List<IContributionItem> items, CheckQuotaRe
 
     // For free / individual / individual_pro users, show an Upgrade Plan row. When the overage row is
     // already showing the upgrade icon directly above, this row uses the blank icon to avoid duplication.
-    if (MenuUtils.shouldShowUpgradePlanRow(plan)) {
+    if (MenuUtils.shouldShowUpgradePlanRow(plan, quotaStatus.canUpgradePlan())) {
       ImageDescriptor upgradePlanIcon = hasNonOrgPremiumQuota ? blankIcon : upgradeIcon;
       items.add(createCommandItem("com.microsoft.copilot.eclipse.commands.upgradeCopilotPlan",
           Messages.menu_quota_upgradePlan, upgradePlanIcon));

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/handlers/ShowStatusBarMenuHandler.java

@@ -339,7 +339,7 @@ private void addCopilotUsageActionTbb(MenuManager menuManager, CheckQuotaResult
 
     // For free / individual / individual_pro users, show an Upgrade Plan row. When the overage row is
     // already showing the upgrade icon directly above, this row uses the blank icon to avoid duplication.
-    if (MenuUtils.shouldShowUpgradePlanRow(plan)) {
+    if (MenuUtils.shouldShowUpgradePlanRow(plan, quotaStatus.canUpgradePlan())) {
       ImageDescriptor upgradePlanIcon = hasNonOrgPremiumQuota ? blankIcon : upgradeIcon;
       MenuActionFactory.createMenuAction(menuManager, Messages.menu_quota_upgradePlan, upgradePlanIcon, handlerService,
           "com.microsoft.copilot.eclipse.commands.upgradeCopilotPlan", true);

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/utils/MenuUtils.java

@@ -126,8 +126,20 @@ public static boolean hasNonOrgPremiumQuota(CheckQuotaResult quotaStatus) {
 
   /**
    * True when the "Upgrade Plan" row should be shown for the given plan.
+   *
+   * <p>{@code canUpgradePlan} is the language server's authoritative signal of whether the user is eligible
+   * to upgrade. When supplied (non-{@code null}) it takes precedence over the plan-based default; when
+   * {@code null} (older language server that does not yet send this field) we fall back to the previous
+   * plan-only heuristic.
+   *
+   * @param plan the user's Copilot plan
+   * @param canUpgradePlan whether the user can upgrade their Copilot plan, or {@code null} when the language
+   *     server did not supply this field
    */
-  public static boolean shouldShowUpgradePlanRow(CopilotPlan plan) {
+  public static boolean shouldShowUpgradePlanRow(CopilotPlan plan, Boolean canUpgradePlan) {
+    if (canUpgradePlan != null) {
+      return canUpgradePlan;
+    }
     return plan == CopilotPlan.free || plan == CopilotPlan.individual || plan == CopilotPlan.individual_pro;
   }