refactor: Update quota usage menu and UI updates (#172)

Author: ethanyhou

com.microsoft.copilot.eclipse.core.test/src/com/microsoft/copilot/eclipse/core/AuthStatusManagerTests.java

@@ -55,7 +55,7 @@ void testSignInConfirm() throws InterruptedException, ExecutionException {
     when(mockResult.getUser()).thenReturn(mockedUser);
     when(mockResult.getStatus()).thenReturn(CopilotStatusResult.OK);
     when(mockConnection.signInConfirm(userCode)).thenReturn(CompletableFuture.completedFuture(mockResult));
-    when(mockConnection.checkQuota()).thenReturn(CompletableFuture.completedFuture(new CheckQuotaResult()));
+    when(mockConnection.checkQuota()).thenReturn(CompletableFuture.completedFuture(CheckQuotaResult.empty()));
 
     CopilotStatusResult result = authStatusManager.signInConfirm(userCode);
 

com.microsoft.copilot.eclipse.core/src/com/microsoft/copilot/eclipse/core/AuthStatusManager.java

@@ -182,7 +182,7 @@ public void setQuotaStatus(CheckQuotaResult checkQuotaResult) {
    */
   public CheckQuotaResult getQuotaStatus() {
     if (this.checkQuotaResult == null) {
-      this.checkQuotaResult = new CheckQuotaResult();
+      this.checkQuotaResult = CheckQuotaResult.empty();
     }
     return this.checkQuotaResult;
   }

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

@@ -3,90 +3,34 @@
 
 package com.microsoft.copilot.eclipse.core.lsp.protocol.quota;
 
-import java.util.Objects;
-
-import org.apache.commons.lang3.builder.ToStringBuilder;
-
 /**
- * Result of the checkQuota request.
+ * Result of the {@code checkQuota} request.
+ *
+ * @param chat chat quota snapshot
+ * @param completions completions quota snapshot
+ * @param premiumInteractions premium interactions quota snapshot
+ * @param resetDate ISO-8601 local date when the monthly allowance resets, or {@code null}
+ * @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
  */
-public class CheckQuotaResult {
-  private Quota chat;
-  private Quota completions;
-  private Quota premiumInteractions;
-  private String resetDate;
-  private CopilotPlan copilotPlan;
-
-  public Quota getChatQuota() {
-    return chat;
-  }
-
-  public void setChatQuota(Quota chat) {
-    this.chat = chat;
-  }
-
-  public Quota getCompletionsQuota() {
-    return completions;
-  }
-
-  public void setCompletionsQuota(Quota completions) {
-    this.completions = completions;
-  }
-
-  public Quota getPremiumInteractionsQuota() {
-    return premiumInteractions;
-  }
-
-  public void setPremiumInteractionsQuota(Quota premiumInteractions) {
-    this.premiumInteractions = premiumInteractions;
-  }
-
-  public String getResetDate() {
-    return resetDate;
-  }
-
-  public void setResetDate(String resetDate) {
-    this.resetDate = resetDate;
-  }
-
-  public CopilotPlan getCopilotPlan() {
-    return copilotPlan;
-  }
-
-  public void setCopilotPlan(CopilotPlan copilotPlan) {
-    this.copilotPlan = copilotPlan;
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hash(chat, completions, copilotPlan, premiumInteractions, resetDate);
-  }
-
-  @Override
-  public boolean equals(Object obj) {
-    if (this == obj) {
-      return true;
-    }
-    if (obj == null) {
-      return false;
-    }
-    if (getClass() != obj.getClass()) {
-      return false;
-    }
-    CheckQuotaResult other = (CheckQuotaResult) obj;
-    return Objects.equals(chat, other.chat) && Objects.equals(completions, other.completions)
-        && copilotPlan == other.copilotPlan && Objects.equals(premiumInteractions, other.premiumInteractions)
-        && Objects.equals(resetDate, other.resetDate);
-  }
-
-  @Override
-  public String toString() {
-    ToStringBuilder builder = new ToStringBuilder(this);
-    builder.append("chat", chat);
-    builder.append("completions", completions);
-    builder.append("premiumInteractions", premiumInteractions);
-    builder.append("resetDate", resetDate);
-    builder.append("copilotPlan", copilotPlan);
-    return builder.toString();
+public record CheckQuotaResult(
+    Quota chat,
+    Quota completions,
+    Quota premiumInteractions,
+    String resetDate,
+    String resetDateUtc,
+    CopilotPlan copilotPlan,
+    boolean tokenBasedBillingEnabled) {
+
+  private static final CheckQuotaResult EMPTY =
+      new CheckQuotaResult(null, null, null, null, null, null, false);
+
+  /**
+   * Returns an empty {@link CheckQuotaResult} used as a placeholder before the language server
+   * supplies real quota data.
+   */
+  public static CheckQuotaResult empty() {
+    return EMPTY;
   }
 }

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

@@ -7,5 +7,5 @@
  * Enum representing the different Copilot plans.
  */
 public enum CopilotPlan {
-  free, individual, individual_pro, business, enterprise
+  free, individual, individual_pro, individual_max, business, enterprise
 }

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

@@ -5,85 +5,61 @@
 
 import java.util.Objects;
 
-import org.apache.commons.lang3.builder.ToStringBuilder;
-
 /**
- * Completions quota information.
+ * Quota information for a single tracked category (chat, completions, or premium interactions).
+ *
+ * <p>Equality intentionally excludes {@link #timeStamp} so that two snapshots with the same
+ * display-meaningful state compare equal even when the language server stamps a different
+ * production time on each refresh.
+ *
+ * @param percentRemaining percentage of the quota remaining; clamped into {@code [0.0, 100.0]} by
+ *     the accessor since the language server may report drift slightly outside that range
+ * @param unlimited whether this category has no monthly limit
+ * @param overagePermitted whether the user has enabled additional paid usage beyond the allowance
+ * @param overageCount additional paid units already consumed, when reported
+ * @param entitlement total monthly allowance, when reported
+ * @param quotaRemaining absolute units remaining in the monthly allowance, when reported
+ * @param timeStamp ISO-8601 timestamp of when the snapshot was produced by the language server;
+ *     not part of {@link #equals(Object)} / {@link #hashCode()}
  */
-public class Quota {
-  private double percentRemaining;
-  private boolean unlimited;
-  private boolean overagePermitted;
-
-  /**
-   * Creates a new CompletionsQuota quota information with default values.
-   */
-  public Quota() {
-    this.percentRemaining = 0.0;
-    this.unlimited = false;
-    this.overagePermitted = false;
-  }
+public record Quota(
+    double percentRemaining,
+    boolean unlimited,
+    boolean overagePermitted,
+    double overageCount,
+    double entitlement,
+    double quotaRemaining,
+    String timeStamp) {
 
   /**
-   * Gets the percentage of the quota remaining within the range of 0.0 to 100.0.
+   * Returns the percentage of the quota remaining, clamped into the {@code [0.0, 100.0]} range.
    */
-  public double getPercentRemaining() {
+  public Quota {
     if (percentRemaining < 0.0) {
-      return 0.0;
+      percentRemaining = 0.0;
     } else if (percentRemaining > 100.0) {
-      return 100.0;
+      percentRemaining = 100.0;
     }
-    return percentRemaining;
-  }
-
-  public void setPercentRemaining(double percentRemaining) {
-    this.percentRemaining = percentRemaining;
-  }
-
-  public boolean isUnlimited() {
-    return unlimited;
-  }
-
-  public void setUnlimited(boolean unlimited) {
-    this.unlimited = unlimited;
-  }
-
-  public boolean isOveragePermitted() {
-    return overagePermitted;
-  }
-
-  public void setOveragePermitted(boolean overagePermitted) {
-    this.overagePermitted = overagePermitted;
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hash(overagePermitted, percentRemaining, unlimited);
   }
 
   @Override
   public boolean equals(Object obj) {
     if (this == obj) {
       return true;
     }
-    if (obj == null) {
+    if (!(obj instanceof Quota other)) {
       return false;
     }
-    if (getClass() != obj.getClass()) {
-      return false;
-    }
-    Quota other = (Quota) obj;
-    return overagePermitted == other.overagePermitted
-        && Double.doubleToLongBits(percentRemaining) == Double.doubleToLongBits(other.percentRemaining)
-        && unlimited == other.unlimited;
+    return Double.compare(percentRemaining, other.percentRemaining) == 0
+        && unlimited == other.unlimited
+        && overagePermitted == other.overagePermitted
+        && Double.compare(overageCount, other.overageCount) == 0
+        && Double.compare(entitlement, other.entitlement) == 0
+        && Double.compare(quotaRemaining, other.quotaRemaining) == 0;
   }
 
   @Override
-  public String toString() {
-    ToStringBuilder builder = new ToStringBuilder(this);
-    builder.append("percentRemaining", percentRemaining);
-    builder.append("unlimited", unlimited);
-    builder.append("overagePermitted", overagePermitted);
-    return builder.toString();
+  public int hashCode() {
+    return Objects.hash(percentRemaining, unlimited, overagePermitted, overageCount, entitlement, quotaRemaining);
   }
 }

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

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+package com.microsoft.copilot.eclipse.core.lsp.protocol.quota;
+
+/**
+ * Snapshot of a single quota bucket (chat, completions, or premium interactions) shipped with
+ * {@code copilot/quotaChange} and {@code copilot/quotaWarning} notifications.
+ *
+ * @param quota total entitlement
+ * @param used computed amount used (entitlement * (1 - percentRemaining / 100))
+ * @param percentRemaining percentage of the quota remaining (0-100)
+ * @param overageUsed overage amount consumed
+ * @param overageEnabled whether overages are permitted
+ * @param resetDate ISO 8601 timestamp when the quota resets, or empty when unknown
+ * @param unlimited true when the quota is unlimited
+ */
+public record QuotaSnapshotParams(double quota, double used, double percentRemaining, double overageUsed,
+    boolean overageEnabled, String resetDate, boolean unlimited) {
+}

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

@@ -38,6 +38,7 @@
 import com.microsoft.copilot.eclipse.ui.chat.services.TodoListService;
 import com.microsoft.copilot.eclipse.ui.i18n.Messages;
 import com.microsoft.copilot.eclipse.ui.swt.CssConstants;
+import com.microsoft.copilot.eclipse.ui.utils.MenuUtils;
 import com.microsoft.copilot.eclipse.ui.utils.SwtUtils;
 
 /**
@@ -218,13 +219,13 @@ public void processTurnEvent(ChatProgressValue value) {
       if (StringUtils.isNotEmpty(errMsg)) {
         // TODO: remove this error message replacement if statement when the CLS side warn message is aligned.
         if (value.getCode() == 402) {
-          CopilotPlan userPlan = this.serviceManager.getAuthStatusManager().getQuotaStatus().getCopilotPlan();
+          CopilotPlan userPlan = this.serviceManager.getAuthStatusManager().getQuotaStatus().copilotPlan();
           CopilotModel fallbackModel = this.serviceManager.getModelService().getFallbackModel();
           String fallbackModelName = fallbackModel != null ? fallbackModel.getModelName()
               : Messages.chat_noQuotaView_fallbackModel;
 
-          if (userPlan == CopilotPlan.individual || userPlan == CopilotPlan.individual_pro) {
-            // Pro and Pro+ message
+          if (MenuUtils.isCfiPlan(userPlan)) {
+            // Pro, Pro+ and Max message
             errMsg = String.format(Messages.chat_noQuotaView_proProplusWarnMsg, fallbackModelName);
           } else if (userPlan == CopilotPlan.business || userPlan == CopilotPlan.enterprise) {
             // CE and CB message
@@ -235,7 +236,7 @@ public void processTurnEvent(ChatProgressValue value) {
         renderWarnMessageWithUpgradePlanButton(errMsg, value.getCode());
 
         if (value.getCode() == 402
-            && this.serviceManager.getAuthStatusManager().getQuotaStatus().getCopilotPlan() != CopilotPlan.free) {
+            && this.serviceManager.getAuthStatusManager().getQuotaStatus().copilotPlan() != CopilotPlan.free) {
           this.serviceManager.getModelService().setFallBackModelAsActiveModel();
           this.serviceManager.getAuthStatusManager().checkQuota();
 

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

@@ -12,6 +12,7 @@
 import org.eclipse.swt.widgets.Composite;
 import org.eclipse.swt.widgets.Label;
 
+import com.microsoft.copilot.eclipse.core.CopilotCore;
 import com.microsoft.copilot.eclipse.ui.chat.services.AvatarService;
 import com.microsoft.copilot.eclipse.ui.chat.services.ChatServiceManager;
 import com.microsoft.copilot.eclipse.ui.i18n.Messages;
@@ -72,8 +73,19 @@ public void renderModelInfo(String modelName, double billingMultiplier) {
         }
         if (StringUtils.isNotBlank(modelName)) {
           Label modelInfoLabel = new Label(footer, SWT.NONE);
-          String formattedMultiplier = ModelUtils.formatBillingMultiplier(billingMultiplier);
-          String displayText = String.format("%s - %s", modelName, formattedMultiplier);
+          // When token-based billing is enabled on the language server, the per-turn billing
+          // multiplier is no longer a meaningful price signal, so render the model name on its
+          // own. Fall back to the legacy "{model} - {multiplier}" format otherwise.
+          boolean tbbEnabled = CopilotCore.getPlugin().getAuthStatusManager()
+              .getQuotaStatus().tokenBasedBillingEnabled();
+          String displayText;
+          if (tbbEnabled) {
+            displayText = modelName;
+          } else {
+            // TODO: Remove this legacy fallback after TBB is officially released.
+            String formattedMultiplier = ModelUtils.formatBillingMultiplier(billingMultiplier);
+            displayText = String.format("%s - %s", modelName, formattedMultiplier);
+          }
           modelInfoLabel.setText(displayText);
           GridData labelGridData = new GridData(SWT.RIGHT, SWT.CENTER, true, false);
           modelInfoLabel.setLayoutData(labelGridData);

com.microsoft.copilot.eclipse.ui/src/com/microsoft/copilot/eclipse/ui/chat/services/ModelService.java

@@ -424,7 +424,7 @@ public void bindModelPicker(final DropdownButton picker) {
       }, (Map<String, CopilotModel> modelMap) -> {
         if (!picker.isDisposed()) {
           boolean showAddPremiumModelOption = this.authStatusManager.getQuotaStatus()
-              .getCopilotPlan() == CopilotPlan.free;
+              .copilotPlan() == CopilotPlan.free;
           // TODO: need to remove this logic after group policy is available
           FeatureFlags flags = CopilotCore.getPlugin().getFeatureFlags();
           boolean showByokManageOption = flags == null || flags.isByokEnabled();