feat: Implement TBB quota warning notifications and related UI components

Author: ethanyhou

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

@@ -10,12 +10,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;
@@ -28,10 +30,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)
@@ -48,6 +52,9 @@ class CopilotLanguageClientTests {
   @Mock
   private IReferencedFileService fileService;
 
+  @Mock
+  private IEventBroker eventBroker;
+
   @BeforeEach
   void setUp() {
     client = new CopilotLanguageClient();
@@ -131,4 +138,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

@@ -30,6 +30,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.
    */
@@ -160,4 +165,9 @@ public class CopilotEventConstants {
    * Event when a rate limit warning is received from the language server.
    */
   public static final String TOPIC_RATE_LIMIT_WARNING = TOPIC_CHAT + "RATE_LIMIT_WARNING";
+
+  /**
+   * 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

@@ -64,6 +64,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;
 
@@ -332,6 +333,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

@@ -14,6 +14,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;
 
@@ -41,6 +43,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;
   }
@@ -59,7 +83,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
@@ -75,7 +100,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);
   }
 
@@ -85,6 +113,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

@@ -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/IntervalQuota.java

@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+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

@@ -14,6 +14,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.
@@ -56,9 +59,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
@@ -73,8 +109,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;
   }
 
@@ -84,6 +121,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/QuotaChangeParams.java

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+package com.microsoft.copilot.eclipse.core.lsp.protocol.quota;
+
+import com.google.gson.annotations.SerializedName;
+
+/**
+ * Parameters for the {@code copilot/quotaChange} notification, sent by the language server
+ * whenever the user's quota usage changes.
+ *
+ * @param chat current chat quota snapshot, when available
+ * @param completions current completions quota snapshot, when available
+ * @param premiumInteractions current premium interactions quota snapshot, when available
+ * @param copilotPlan the user's Copilot plan (e.g. free, individual, individual_pro, individual_max,
+ *     business, enterprise)
+ */
+public record QuotaChangeParams(QuotaSnapshotParams chat, QuotaSnapshotParams completions,
+    @SerializedName("premium_interactions") QuotaSnapshotParams premiumInteractions, String copilotPlan) {
+}

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

@@ -0,0 +1,11 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+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.core/src/com/microsoft/copilot/eclipse/core/lsp/protocol/quota/QuotaWarningParams.java

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+package com.microsoft.copilot.eclipse.core.lsp.protocol.quota;
+
+import com.google.gson.annotations.SerializedName;
+
+/**
+ * Parameters for the {@code copilot/quotaWarning} notification, sent by the language server when
+ * the user crosses a quota usage threshold or starts consuming overages.
+ *
+ * @param title warning title (e.g. "Copilot Quota Usage Alert")
+ * @param message human-readable warning message
+ * @param severity severity level, either {@code "warning"} or {@code "info"}
+ * @param chat current chat quota snapshot, when available
+ * @param completions current completions quota snapshot, when available
+ * @param premiumInteractions current premium interactions quota snapshot, when available
+ * @param copilotPlan the user's Copilot plan
+ */
+public record QuotaWarningParams(String title, String message, String severity, QuotaSnapshotParams chat,
+    QuotaSnapshotParams completions,
+    @SerializedName("premium_interactions") QuotaSnapshotParams premiumInteractions, String copilotPlan) {
+}