Add usage limits UI for copilot

workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts

@@ -44,6 +44,7 @@ import {
     UIChatMessage,
     CheckpointInfo,
     AbortAIGenerationRequest,
+    UsageResponse,
 } from "./interfaces";
 
 export interface AIPanelAPI {
@@ -107,4 +108,5 @@ export interface AIPanelAPI {
     clearChat: () => Promise<void>;
     updateChatMessage: (params: UpdateChatMessageRequest) => Promise<void>;
     getActiveTempDir: () => Promise<string>;
+    getUsage: () => Promise<UsageResponse | undefined>;
 }

workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/interfaces.ts

@@ -486,3 +486,8 @@ export interface AbortAIGenerationRequest {
     /** Thread identifier (defaults to 'default') */
     threadId?: string;
 }
+
+export interface UsageResponse {
+    remainingUsagePercentage: number;
+    resetsIn: number; // in seconds
+}

workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts

@@ -46,6 +46,7 @@ import {
     UIChatMessage,
     CheckpointInfo,
     AbortAIGenerationRequest,
+    UsageResponse,
 } from "./interfaces";
 import { RequestType, NotificationType } from "vscode-messenger-common";
 
@@ -96,3 +97,4 @@ export const restoreCheckpoint: RequestType<RestoreCheckpointRequest, void> = {
 export const clearChat: RequestType<void, void> = { method: `${_preFix}/clearChat` };
 export const updateChatMessage: RequestType<UpdateChatMessageRequest, void> = { method: `${_preFix}/updateChatMessage` };
 export const getActiveTempDir: RequestType<void, string> = { method: `${_preFix}/getActiveTempDir` };
+export const getUsage: RequestType<void, UsageResponse | undefined> = { method: `${_preFix}/getUsage` };

workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-client.ts

@@ -134,7 +134,7 @@ export async function fetchWithAuth(input: string | URL | Request, options: Requ
         // Handle usage limit exceeded
         if (response.status === 429) {
             console.log("Usage limit exceeded (429)");
-            const error = new Error("Usage limit exceeded. Please try again later.");
+            const error = new Error("Usage limit exceeded.");
             error.name = "UsageLimitError";
             (error as any).statusCode = 429;
             throw error;

workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-utils.ts

@@ -313,7 +313,7 @@ export function getErrorMessage(error: unknown): string {
     if (error instanceof Error) {
         // Standard Error objects have a .message property
         if (error.name === "UsageLimitError") {
-            return "Usage limit exceeded. Please try again later.";
+            return "Usage limit exceeded.";
         }
         if (error.name === "AI_RetryError") {
             return "An error occured connecting with the AI service. Please try again later.";
@@ -333,7 +333,7 @@ export function getErrorMessage(error: unknown): string {
     ) {
         // Check if it has a statusCode property indicating 429
         if ("statusCode" in error && (error as any).statusCode === 429) {
-            return "Usage limit exceeded. Please try again later.";
+            return "Usage limit exceeded.";
         }
         return (error as { message: string }).message;
     }

workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts

@@ -84,7 +84,8 @@ import {
     TaskDeclineRequest,
     updateChatMessage,
     UpdateChatMessageRequest,
-    updateRequirementSpecification
+    updateRequirementSpecification,
+    getUsage
 } from "@wso2/ballerina-core";
 import { Messenger } from "vscode-messenger";
 import { AiPanelRpcManager } from "./rpc-manager";
@@ -137,4 +138,5 @@ export function registerAiPanelRpcHandlers(messenger: Messenger) {
     messenger.onRequest(clearChat, () => rpcManger.clearChat());
     messenger.onRequest(updateChatMessage, (args: UpdateChatMessageRequest) => rpcManger.updateChatMessage(args));
     messenger.onRequest(getActiveTempDir, () => rpcManger.getActiveTempDir());
+    messenger.onRequest(getUsage, () => rpcManger.getUsage());
 }

workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts

@@ -41,7 +41,8 @@ import {
     SubmitFeedbackRequest,
     TestGenerationMentions,
     UIChatMessage,
-    UpdateChatMessageRequest
+    UpdateChatMessageRequest,
+    UsageResponse
 } from "@wso2/ballerina-core";
 import * as fs from 'fs';
 import path from "path";
@@ -65,12 +66,14 @@ import { refreshDataMapper } from "../data-mapper/utils";
 import {
     TEST_DIR_NAME
 } from "./constants";
+import { fetchWithAuth } from "../../features/ai/utils/ai-client";
+import { BACKEND_URL } from "../../features/ai/utils";
 import { addToIntegration, cleanDiagnosticMessages, searchDocumentation } from "./utils";
 
 import { onHideReviewActions } from '@wso2/ballerina-core';
 import { createExecutionContextFromStateMachine, createExecutorConfig, generateAgent } from '../../features/ai/agent/index';
 import { integrateCodeToWorkspace } from "../../features/ai/agent/utils";
-import { WI_EXTENSION_ID } from "../../features/ai/constants";
+import { LLM_API_BASE_PATH, WI_EXTENSION_ID } from "../../features/ai/constants";
 import { ContextTypesExecutor } from '../../features/ai/executors/datamapper/ContextTypesExecutor';
 import { FunctionMappingExecutor } from '../../features/ai/executors/datamapper/FunctionMappingExecutor';
 import { InlineMappingExecutor } from '../../features/ai/executors/datamapper/InlineMappingExecutor';
@@ -699,4 +702,24 @@ export class AiPanelRpcManager implements AIPanelAPI {
         console.log(">>> active temp project path", projectPath);
         return projectPath;
     }
+
+    async getUsage(): Promise<UsageResponse | undefined> {
+        const loginMethod = await getLoginMethod();
+        if (loginMethod !== LoginMethod.BI_INTEL) {
+            return undefined;
+        }
+        try {
+            const url = BACKEND_URL + LLM_API_BASE_PATH + "/usage";
+            const response = await fetchWithAuth(url, { method: "GET" });
+            if (response && response.ok) {
+                const data = await response.json();
+                return data as UsageResponse;
+            }
+            console.error("Failed to fetch usage: ", response?.status, response?.statusText);
+            return undefined;
+        } catch (error) {
+            console.error("Failed to fetch usage:", error);
+            return undefined;
+        }
+    }
 }

workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiMachine.ts

@@ -33,8 +33,6 @@ import {
 import * as vscode from 'vscode';
 import { notifyAiPromptUpdated } from '../../RPCLayer';
 
-export const USER_CHECK_BACKEND_URL = '/user/usage';
-
 export const openAIWebview = (defaultprompt?: AIPanelPrompt) => {
     extension.aiChatDefaultPrompt = defaultprompt;
     if (!AiPanelWebview.currentPanel) {

workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts

@@ -47,6 +47,7 @@ import {
     TestGenerationMentions,
     UIChatMessage,
     UpdateChatMessageRequest,
+    UsageResponse,
     abortAIGeneration,
     acceptChanges,
     addFilesToProject,
@@ -92,7 +93,8 @@ import {
     showSignInAlert,
     submitFeedback,
     updateChatMessage,
-    updateRequirementSpecification
+    updateRequirementSpecification,
+    getUsage
 } from "@wso2/ballerina-core";
 import { HOST_EXTENSION } from "vscode-messenger-common";
 import { Messenger } from "vscode-messenger-webview";
@@ -287,4 +289,8 @@ export class AiPanelRpcClient implements AIPanelAPI {
     getActiveTempDir(): Promise<string> {
         return this._messenger.sendRequest(getActiveTempDir, HOST_EXTENSION);
     }
+
+    getUsage(): Promise<UsageResponse | undefined> {
+        return this._messenger.sendRequest(getUsage, HOST_EXTENSION);
+    }
 }

workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/index.tsx

@@ -134,6 +134,7 @@ type FooterProps = {
     onChangeAgentMode?: (mode: AgentMode) => void;
     isAutoApproveEnabled?: boolean;
     onDisableAutoApprove?: () => void;
+    disabled?: boolean;
 };
 
 const Footer: React.FC<FooterProps> = ({
@@ -151,6 +152,7 @@ const Footer: React.FC<FooterProps> = ({
     onChangeAgentMode,
     isAutoApproveEnabled,
     onDisableAutoApprove,
+    disabled,
 }) => {
     const [generatingText, setGeneratingText] = useState("Generating.");
 
@@ -200,6 +202,7 @@ const Footer: React.FC<FooterProps> = ({
                 onChangeAgentMode={onChangeAgentMode}
                 isAutoApproveEnabled={isAutoApproveEnabled}
                 onDisableAutoApprove={onDisableAutoApprove}
+                disabled={disabled}
             />
         </FooterContainer>
     );

workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx

@@ -51,7 +51,7 @@ import RoleContainer from "../RoleContainter";
 import CheckpointSeparator from "../CheckpointSeparator";
 import { Attachment, AttachmentStatus, TaskApprovalRequest } from "@wso2/ballerina-core";
 
-import { AIChatView, Header, HeaderButtons, ChatMessage, Badge, ApprovalOverlay, OverlayMessage } from "../../styles";
+import { AIChatView, Header, HeaderButtons, ChatMessage, Badge, ResetsInBadge, ApprovalOverlay, OverlayMessage } from "../../styles";
 import ReferenceDropdown from "../ReferenceDropdown";
 import { VSCodeButton } from "@vscode/webview-ui-toolkit/react";
 import MarkdownRenderer from "../MarkdownRenderer";
@@ -89,6 +89,8 @@ const DRIFT_CHECK_ERROR = "Failed to check drift between the code and the docume
 const GENERATE_CODE_AGAINST_THE_PROVIDED_REQUIREMENTS = "Generate code based on the following requirements: ";
 const GENERATE_CODE_AGAINST_THE_PROVIDED_REQUIREMENTS_TRIMMED = GENERATE_CODE_AGAINST_THE_PROVIDED_REQUIREMENTS.trim();
 
+const USAGE_EXCEEDED_THRESHOLD_PERCENT = 3;
+
 /**
  * Formats a file path into a user-friendly display name
  * - Removes .bal extension
@@ -156,6 +158,9 @@ const AIChat: React.FC = () => {
     const [currentFileArray, setCurrentFileArray] = useState<SourceFile[]>([]);
     const [codeContext, setCodeContext] = useState<CodeContext | undefined>(undefined);
 
+    const [usage, setUsage] = useState<{ remainingUsagePercentage: number; resetsIn: number } | null>(null);
+    const [isUsageExceeded, setIsUsageExceeded] = useState(false);
+
     //TODO: Need a better way of storing data related to last generation to be in the repair state.
     const currentDiagnosticsRef = useRef<DiagnosticEntry[]>([]);
     const functionsRef = useRef<any>([]);
@@ -226,6 +231,46 @@ const AIChat: React.FC = () => {
     }, []);
     /* REFACTORED CODE END [2] */
 
+    const formatResetsIn = (seconds: number): string => {
+        const days = Math.floor(seconds / 86400);
+        if (days >= 1) return `${days} day${days > 1 ? 's' : ''}`;
+        const hours = Math.floor(seconds / 3600);
+        if (hours >= 1) return `${hours} hour${hours > 1 ? 's' : ''}`;
+        const mins = Math.floor(seconds / 60);
+        return `${mins} min${mins > 1 ? 's' : ''}`;
+    };
+
+    const formatResetsInExact = (seconds: number): string => {
+        const days = Math.floor(seconds / 86400);
+        const hours = Math.floor((seconds % 86400) / 3600);
+        const mins = Math.floor((seconds % 3600) / 60);
+        const parts: string[] = [];
+        if (days > 0) parts.push(`${days} day${days > 1 ? 's' : ''}`);
+        if (hours > 0) parts.push(`${hours} hour${hours > 1 ? 's' : ''}`);
+        if (mins > 0) parts.push(`${mins} minute${mins > 1 ? 's' : ''}`);
+        return parts.length > 0 ? parts.join(', ') : 'less than a minute';
+    };
+
+    const fetchUsage = async () => {
+        try {
+            const result = await rpcClient.getAiPanelRpcClient().getUsage();
+            if (result) {
+                setUsage(result);
+                setIsUsageExceeded(result.resetsIn !== -1 && result.remainingUsagePercentage < USAGE_EXCEEDED_THRESHOLD_PERCENT);
+            } else {
+                setUsage(null);
+                setIsUsageExceeded(false);
+            }
+        } catch (e) {
+            console.error("Failed to fetch usage:", e);
+            // Reset on error to avoid permanently blocking the user on transient failures
+            setUsage(null);
+            setIsUsageExceeded(false);
+        }
+    };
+
+    useEffect(() => { fetchUsage(); }, []);
+
     const handleCheckpointRestore = async (checkpointId: string) => {
         try {
             // Call backend to restore checkpoint (files + chat history)
@@ -730,6 +775,7 @@ const AIChat: React.FC = () => {
             console.log("Received stop signal");
             setIsCodeLoading(false);
             setIsLoading(false);
+            fetchUsage();
         } else if (type === "abort") {
             console.log("Received abort signal");
             const interruptedMessage = "\n\n*[Request interrupted by user]*";
@@ -1407,9 +1453,19 @@ const AIChat: React.FC = () => {
                     )}
                     <Header>
                         <Badge>
-                            Remaining Free Usage: {"Unlimited"}
-                            <br />
-                            {/* <ResetsInBadge>{`Resets in: 30 days`}</ResetsInBadge> */}
+                            {usage ? (
+                                <>
+                                    Remaining Usage: {usage.resetsIn === -1 ? "Unlimited" : (isUsageExceeded ? "Exceeded" : `${Math.round(usage.remainingUsagePercentage)}%`)}
+                                    {usage.resetsIn !== -1 && (
+                                        <>
+                                            <br />
+                                            <ResetsInBadge title={formatResetsInExact(usage.resetsIn)}>{`Resets in: ${formatResetsIn(usage.resetsIn)}`}</ResetsInBadge>
+                                        </>
+                                    )}
+                                </>
+                            ) : (
+                                "Remaining Usage: N/A"
+                            )}
                         </Badge>
                         <HeaderButtons>
                             <Button
@@ -1784,6 +1840,7 @@ const AIChat: React.FC = () => {
                             onChangeAgentMode={isPlanModeFeatureEnabled ? handleChangeAgentMode : undefined}
                             isAutoApproveEnabled={isAutoApproveEnabled}
                             onDisableAutoApprove={handleToggleAutoApprove}
+                            disabled={isUsageExceeded}
                         />
                     )}
                 </AIChatView>

workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChatInput/StyledInput.tsx

@@ -121,10 +121,11 @@ interface StyledInputProps {
     onBlur: (e: React.FocusEvent<HTMLDivElement>) => void;
     placeholder: string;
     onPostDOMUpdate?: () => void;
+    disabled?: boolean;
 }
 
 export const StyledInputComponent = forwardRef<StyledInputRef, StyledInputProps>(
-    ({ value, onChange, onKeyDown, onBlur, placeholder, onPostDOMUpdate }, ref) => {
+    ({ value, onChange, onKeyDown, onBlur, placeholder, onPostDOMUpdate, disabled }, ref) => {
         const [internalContent, setInternalContent] = useState<string>(value.text);
         const divRef = useRef<HTMLDivElement>(null);
 
@@ -395,7 +396,7 @@ export const StyledInputComponent = forwardRef<StyledInputRef, StyledInputProps>
         return (
             <StyledInput
                 ref={divRef}
-                contentEditable
+                contentEditable={!disabled}
                 spellCheck="true"
                 onInput={handleInput}
                 onKeyDown={handleKeyDown}

workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChatInput/index.tsx

@@ -132,11 +132,12 @@ interface AIChatInputProps {
     onChangeAgentMode?: (mode: AgentMode) => void;
     isAutoApproveEnabled?: boolean;
     onDisableAutoApprove?: () => void;
+    disabled?: boolean;
 }
 
 const AIChatInput = forwardRef<AIChatInputRef, AIChatInputProps>(
     ({ initialCommandTemplate, tagOptions, attachmentOptions, placeholder, onSend, onStop, isLoading,
-       agentMode = AgentMode.Edit, onChangeAgentMode, isAutoApproveEnabled = false, onDisableAutoApprove }, ref) => {
+       agentMode = AgentMode.Edit, onChangeAgentMode, isAutoApproveEnabled = false, onDisableAutoApprove, disabled }, ref) => {
         const [inputValue, setInputValue] = useState<{
             text: string;
             [key: string]: any;
@@ -540,8 +541,9 @@ const AIChatInput = forwardRef<AIChatInputRef, AIChatInputProps>(
                             onChange={setInputValue}
                             onKeyDown={handleKeyDown}
                             onBlur={() => completeSuggestionSelection()}
-                            placeholder={placeholder}
+                            placeholder={disabled ? "Usage limit exceeded" : placeholder}
                             onPostDOMUpdate={executeOnPostDOMUpdate}
+                            disabled={disabled}
                         />
                         {/* Attachments Display */}
                         {attachments.length > 0 && (
@@ -594,7 +596,7 @@ const AIChatInput = forwardRef<AIChatInputRef, AIChatInputProps>(
                             <div>
                                 <ActionButton
                                     title={isLoading ? "Stop (Escape)" : "Send (Enter)"}
-                                    disabled={inputValue.text.trim() === "" && !isLoading}
+                                    disabled={(inputValue.text.trim() === "" && !isLoading) || disabled}
                                     onClick={isLoading ? handleStop : handleSend}
                                 >
                                     <span