feat: [New billing] Add warning notification for copilot usage limit. (#1621)

Author: ethanyhou

com.microsoft.copilot.eclipse.core.test/src/com/microsoft/copilot/eclipse/core/lsp/CopilotLanguageClientTests.java

@@ -7,12 +7,14 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 
 import org.eclipse.core.resources.IFile;
+import org.eclipse.e4.core.services.events.IEventBroker;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -25,10 +27,12 @@
 import com.microsoft.copilot.eclipse.core.FeatureFlags;
 import com.microsoft.copilot.eclipse.core.chat.service.IChatServiceManager;
 import com.microsoft.copilot.eclipse.core.chat.service.IReferencedFileService;
+import com.microsoft.copilot.eclipse.core.events.CopilotEventConstants;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.ConversationCapabilities;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.ConversationContextParams;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.CurrentEditorContext;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.DidChangeFeatureFlagsParams;
+import com.microsoft.copilot.eclipse.core.lsp.protocol.quota.QuotaWarningNotification;
 import com.microsoft.copilot.eclipse.core.utils.FileUtils;
 
 @ExtendWith(MockitoExtension.class)
@@ -45,6 +49,9 @@ class CopilotLanguageClientTests {
   @Mock
   private IReferencedFileService fileService;
 
+  @Mock
+  private IEventBroker eventBroker;
+
   @BeforeEach
   void setUp() {
     client = new CopilotLanguageClient();
@@ -128,4 +135,24 @@ void testOnDidChangeFeatureFlagsWithEmptyFeatureFlags() {
       verify(mockFeatureFlags).setByokEnabled(true);
     }
   }
+
+  @Test
+  void testOnQuotaWarning_PostsNotificationToEventBroker() {
+    QuotaWarningNotification notification = new QuotaWarningNotification("Approaching quota", 90.0);
+    setEventBroker(eventBroker);
+
+    client.onQuotaWarning(notification);
+
+    verify(eventBroker).post(CopilotEventConstants.TOPIC_QUOTA_WARNING, notification);
+  }
+
+  private void setEventBroker(IEventBroker broker) {
+    try {
+      Field field = CopilotLanguageClient.class.getDeclaredField("eventBroker");
+      field.setAccessible(true);
+      field.set(client, broker);
+    } catch (ReflectiveOperationException e) {
+      throw new AssertionError("Failed to inject event broker", e);
+    }
+  }
 }

com.microsoft.copilot.eclipse.core/src/com/microsoft/copilot/eclipse/core/events/CopilotEventConstants.java

@@ -27,6 +27,11 @@ public class CopilotEventConstants {
    */
   private static final String TOPIC_MCP = TOPIC_BASE + "MCP/";
 
+  /**
+   * Topic for quota events.
+   */
+  private static final String TOPIC_QUOTA = TOPIC_BASE + "QUOTA/";
+
   /**
    * Topic for Next Edit Suggestion (NES) events.
    */
@@ -152,4 +157,9 @@ public class CopilotEventConstants {
    * Event when NES suggestion is rejected.
    */
   public static final String TOPIC_NES_REJECT_SUGGESTION = TOPIC_NES + "REJECT_SUGGESTION";
+
+  /**
+   * Event when a quota warning notification is received from the language server.
+   */
+  public static final String TOPIC_QUOTA_WARNING = TOPIC_QUOTA + "WARNING";
 }

com.microsoft.copilot.eclipse.core/src/com/microsoft/copilot/eclipse/core/lsp/CopilotLanguageClient.java

@@ -55,6 +55,7 @@
 import com.microsoft.copilot.eclipse.core.lsp.protocol.codingagent.CodingAgentMessageRequestParams;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.codingagent.CodingAgentMessageResult;
 import com.microsoft.copilot.eclipse.core.lsp.protocol.policy.DidChangePolicyParams;
+import com.microsoft.copilot.eclipse.core.lsp.protocol.quota.QuotaWarningNotification;
 import com.microsoft.copilot.eclipse.core.utils.FileUtils;
 import com.microsoft.copilot.eclipse.core.utils.PlatformUtils;
 
@@ -313,6 +314,17 @@ public CompletableFuture<CodingAgentMessageResult> onCodingAgentMessage(CodingAg
     return CompletableFuture.completedFuture(result);
   }
 
+  /**
+   * Notify when a quota warning is received from the language server.
+   */
+  @JsonNotification("copilot/quotaWarning")
+  public void onQuotaWarning(QuotaWarningNotification notification) {
+    CopilotCore.LOGGER.info("Quota warning received: " + notification);
+    if (eventBroker != null) {
+      eventBroker.post(CopilotEventConstants.TOPIC_QUOTA_WARNING, notification);
+    }
+  }
+
   /**
    * Reads the contents and stats of a file given its URI.
    */

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

@@ -11,6 +11,8 @@ public class CheckQuotaResult {
   private Quota chat;
   private Quota completions;
   private Quota premiumInteractions;
+  private IntervalQuota immediateUsageInterval;
+  private IntervalQuota extendedUsageInterval;
   private String resetDate;
   private CopilotPlan copilotPlan;
 
@@ -38,6 +40,28 @@ public void setPremiumInteractionsQuota(Quota premiumInteractions) {
     this.premiumInteractions = premiumInteractions;
   }
 
+  /**
+   * Gets the immediate usage interval quota (for individual plans).
+   */
+  public IntervalQuota getImmediateUsageInterval() {
+    return immediateUsageInterval;
+  }
+
+  public void setImmediateUsageInterval(IntervalQuota immediateUsageInterval) {
+    this.immediateUsageInterval = immediateUsageInterval;
+  }
+
+  /**
+   * Gets the extended usage interval quota (for individual plans).
+   */
+  public IntervalQuota getExtendedUsageInterval() {
+    return extendedUsageInterval;
+  }
+
+  public void setExtendedUsageInterval(IntervalQuota extendedUsageInterval) {
+    this.extendedUsageInterval = extendedUsageInterval;
+  }
+
   public String getResetDate() {
     return resetDate;
   }
@@ -56,7 +80,8 @@ public void setCopilotPlan(CopilotPlan copilotPlan) {
 
   @Override
   public int hashCode() {
-    return Objects.hash(chat, completions, copilotPlan, premiumInteractions, resetDate);
+    return Objects.hash(chat, completions, copilotPlan, extendedUsageInterval,
+        immediateUsageInterval, premiumInteractions, resetDate);
   }
 
   @Override
@@ -72,7 +97,10 @@ public boolean equals(Object obj) {
     }
     CheckQuotaResult other = (CheckQuotaResult) obj;
     return Objects.equals(chat, other.chat) && Objects.equals(completions, other.completions)
-        && copilotPlan == other.copilotPlan && Objects.equals(premiumInteractions, other.premiumInteractions)
+        && copilotPlan == other.copilotPlan
+        && Objects.equals(extendedUsageInterval, other.extendedUsageInterval)
+        && Objects.equals(immediateUsageInterval, other.immediateUsageInterval)
+        && Objects.equals(premiumInteractions, other.premiumInteractions)
         && Objects.equals(resetDate, other.resetDate);
   }
 
@@ -82,6 +110,8 @@ public String toString() {
     builder.append("chat", chat);
     builder.append("completions", completions);
     builder.append("premiumInteractions", premiumInteractions);
+    builder.append("immediateUsageInterval", immediateUsageInterval);
+    builder.append("extendedUsageInterval", extendedUsageInterval);
     builder.append("resetDate", resetDate);
     builder.append("copilotPlan", copilotPlan);
     return builder.toString();

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

@@ -4,5 +4,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/IntervalQuota.java

@@ -0,0 +1,14 @@
+package com.microsoft.copilot.eclipse.core.lsp.protocol.quota;
+
+/**
+ * Interval-based quota information, used for immediateUsageInterval and extendedUsageInterval.
+ */
+public record IntervalQuota(
+    double percentRemaining,
+    boolean unlimited,
+    boolean overagePermitted,
+    Integer entitlement,
+    Integer quotaRemaining,
+    String timeStamp,
+    String resetAt) {
+}

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

@@ -11,6 +11,9 @@ public class Quota {
   private double percentRemaining;
   private boolean unlimited;
   private boolean overagePermitted;
+  private Integer entitlement;
+  private Integer quotaRemaining;
+  private String timeStamp;
 
   /**
    * Creates a new CompletionsQuota quota information with default values.
@@ -53,9 +56,42 @@ public void setOveragePermitted(boolean overagePermitted) {
     this.overagePermitted = overagePermitted;
   }
 
+  /**
+   * Gets the total entitlement (quota limit).
+   */
+  public Integer getEntitlement() {
+    return entitlement;
+  }
+
+  public void setEntitlement(Integer entitlement) {
+    this.entitlement = entitlement;
+  }
+
+  /**
+   * Gets the remaining quota count.
+   */
+  public Integer getQuotaRemaining() {
+    return quotaRemaining;
+  }
+
+  public void setQuotaRemaining(Integer quotaRemaining) {
+    this.quotaRemaining = quotaRemaining;
+  }
+
+  /**
+   * Gets the timestamp of the quota snapshot.
+   */
+  public String getTimeStamp() {
+    return timeStamp;
+  }
+
+  public void setTimeStamp(String timeStamp) {
+    this.timeStamp = timeStamp;
+  }
+
   @Override
   public int hashCode() {
-    return Objects.hash(overagePermitted, percentRemaining, unlimited);
+    return Objects.hash(entitlement, overagePermitted, percentRemaining, quotaRemaining, timeStamp, unlimited);
   }
 
   @Override
@@ -70,8 +106,9 @@ public boolean equals(Object obj) {
       return false;
     }
     Quota other = (Quota) obj;
-    return overagePermitted == other.overagePermitted
+    return Objects.equals(entitlement, other.entitlement) && overagePermitted == other.overagePermitted
         && Double.doubleToLongBits(percentRemaining) == Double.doubleToLongBits(other.percentRemaining)
+        && Objects.equals(quotaRemaining, other.quotaRemaining) && Objects.equals(timeStamp, other.timeStamp)
         && unlimited == other.unlimited;
   }
 
@@ -81,6 +118,9 @@ public String toString() {
     builder.append("percentRemaining", percentRemaining);
     builder.append("unlimited", unlimited);
     builder.append("overagePermitted", overagePermitted);
+    builder.append("entitlement", entitlement);
+    builder.append("quotaRemaining", quotaRemaining);
+    builder.append("timeStamp", timeStamp);
     return builder.toString();
   }
 }

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

@@ -0,0 +1,8 @@
+package com.microsoft.copilot.eclipse.core.lsp.protocol.quota;
+
+/**
+ * Parameters for the "copilot/quotaWarning" notification. Sent by the language server when the user's AI quota exceeds
+ * the warning threshold.
+ */
+public record QuotaWarningNotification(String message, double percentUsed) {
+}

com.microsoft.copilot.eclipse.ui.test/src/com/microsoft/copilot/eclipse/ui/i18n/MessagesTest.java

@@ -11,5 +11,8 @@ void testMessagesInitialization() {
     // Ensure that the static fields are initialized
     assertNotNull(Messages.menu_signToGitHub);
     assertNotNull(Messages.menu_signOutOfGitHub);
+    assertNotNull(Messages.quotaWarning_title);
+    assertNotNull(Messages.quotaWarning_closeButton);
+    assertNotNull(Messages.quotaWarning_increaseBudgetButton);
   }
 }