From b15b5e8821b6415193f22b6d418535df695e427a Mon Sep 17 00:00:00 2001 From: Anjana Supun Date: Wed, 18 Feb 2026 10:21:28 +0530 Subject: [PATCH 01/19] Usage limit bug fixes --- .../src/rpc-types/ai-panel/index.ts | 2 + .../src/rpc-types/ai-panel/interfaces.ts | 5 ++ .../src/rpc-types/ai-panel/rpc-type.ts | 2 + .../src/features/ai/utils/ai-client.ts | 2 +- .../src/features/ai/utils/ai-utils.ts | 4 +- .../src/rpc-managers/ai-panel/rpc-handler.ts | 4 +- .../src/rpc-managers/ai-panel/rpc-manager.ts | 26 +++++++- .../src/views/ai-panel/aiMachine.ts | 2 - .../src/rpc-clients/ai-panel/rpc-client.ts | 8 ++- .../components/AIChat/Footer/index.tsx | 3 + .../views/AIPanel/components/AIChat/index.tsx | 60 +++++++++++++++++-- .../components/AIChatInput/StyledInput.tsx | 5 +- .../AIPanel/components/AIChatInput/index.tsx | 8 ++- 13 files changed, 113 insertions(+), 18 deletions(-) diff --git a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts b/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts index df5f0f1a6df..4f4bc4c922b 100644 --- a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts +++ b/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; updateChatMessage: (params: UpdateChatMessageRequest) => Promise; getActiveTempDir: () => Promise; + getUsage: () => Promise; } diff --git a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/interfaces.ts b/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/interfaces.ts index f82730c9259..3b4bd0fdf5d 100644 --- a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/interfaces.ts +++ b/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 +} diff --git a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts b/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts index 4bcbe84887e..55e18d09818 100644 --- a/workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts +++ b/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 = { export const clearChat: RequestType = { method: `${_preFix}/clearChat` }; export const updateChatMessage: RequestType = { method: `${_preFix}/updateChatMessage` }; export const getActiveTempDir: RequestType = { method: `${_preFix}/getActiveTempDir` }; +export const getUsage: RequestType = { method: `${_preFix}/getUsage` }; diff --git a/workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-client.ts b/workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-client.ts index ce6de54d35c..53ae9b0393b 100644 --- a/workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-client.ts +++ b/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; diff --git a/workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-utils.ts b/workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-utils.ts index 83f59b21e9f..7227d2ea4a0 100644 --- a/workspaces/ballerina/ballerina-extension/src/features/ai/utils/ai-utils.ts +++ b/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; } diff --git a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts index c0bbb2d4c4e..7ce1c8e21bd 100644 --- a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-handler.ts +++ b/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()); } diff --git a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts b/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts index 5e42c9633e2..eaa70ca345b 100644 --- a/workspaces/ballerina/ballerina-extension/src/rpc-managers/ai-panel/rpc-manager.ts +++ b/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,23 @@ export class AiPanelRpcManager implements AIPanelAPI { console.log(">>> active temp project path", projectPath); return projectPath; } + + async getUsage(): Promise { + 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; + } + return undefined; + } catch (error) { + console.error("Failed to fetch usage:", error); + return undefined; + } + } } diff --git a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiMachine.ts b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiMachine.ts index 31468e1723f..6294ad1b483 100644 --- a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiMachine.ts +++ b/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) { diff --git a/workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts b/workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts index fbf8df00ed2..9c641875b49 100644 --- a/workspaces/ballerina/ballerina-rpc-client/src/rpc-clients/ai-panel/rpc-client.ts +++ b/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 { return this._messenger.sendRequest(getActiveTempDir, HOST_EXTENSION); } + + getUsage(): Promise { + return this._messenger.sendRequest(getUsage, HOST_EXTENSION); + } } diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/index.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/index.tsx index e0e7f64c543..6582e568287 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/Footer/index.tsx +++ b/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 = ({ @@ -151,6 +152,7 @@ const Footer: React.FC = ({ onChangeAgentMode, isAutoApproveEnabled, onDisableAutoApprove, + disabled, }) => { const [generatingText, setGeneratingText] = useState("Generating."); @@ -200,6 +202,7 @@ const Footer: React.FC = ({ onChangeAgentMode={onChangeAgentMode} isAutoApproveEnabled={isAutoApproveEnabled} onDisableAutoApprove={onDisableAutoApprove} + disabled={disabled} /> ); diff --git a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx b/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx index ead4f6b2d0c..7b22f8c573b 100644 --- a/workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx +++ b/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"; @@ -156,6 +156,9 @@ const AIChat: React.FC = () => { const [currentFileArray, setCurrentFileArray] = useState([]); const [codeContext, setCodeContext] = useState(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([]); const functionsRef = useRef([]); @@ -226,6 +229,43 @@ 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 < 3); + } else { + setUsage(null); + setIsUsageExceeded(false); + } + } catch (e) { + console.error("Failed to fetch usage:", e); + } + }; + + useEffect(() => { fetchUsage(); }, []); + const handleCheckpointRestore = async (checkpointId: string) => { try { // Call backend to restore checkpoint (files + chat history) @@ -730,6 +770,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 +1448,19 @@ const AIChat: React.FC = () => { )}
- Remaining Free Usage: {"Unlimited"} -
- {/* {`Resets in: 30 days`} */} + {usage ? ( + <> + Remaining Usage: {usage.resetsIn === -1 ? "Unlimited" : (isUsageExceeded ? "Exceeded" : `${Math.round(usage.remainingUsagePercentage)}%`)} + {usage.resetsIn !== -1 && ( + <> +
+ {`Resets in: ${formatResetsIn(usage.resetsIn)}`} + + )} + + ) : ( + "Remaining Usage: N/A" + )}