Add auto-approved cost limits (#6484)

Co-authored-by: Daniel Riccio <[email protected]>

Author: hassoncs

packages/types/src/global-settings.ts

@@ -68,6 +68,7 @@ export const globalSettingsSchema = z.object({
 	commandTimeoutAllowlist: z.array(z.string()).optional(),
 	preventCompletionWithOpenTodos: z.boolean().optional(),
 	allowedMaxRequests: z.number().nullish(),
+	allowedMaxCost: z.number().nullish(),
 	autoCondenseContext: z.boolean().optional(),
 	autoCondenseContextPercent: z.number().optional(),
 	maxConcurrentFileReads: z.number().optional(),

src/core/task/AutoApprovalHandler.ts

@@ -0,0 +1,144 @@
+import { GlobalState, ClineMessage, ClineAsk } from "@roo-code/types"
+import { getApiMetrics } from "../../shared/getApiMetrics"
+import { ClineAskResponse } from "../../shared/WebviewMessage"
+
+export interface AutoApprovalResult {
+	shouldProceed: boolean
+	requiresApproval: boolean
+	approvalType?: "requests" | "cost"
+	approvalCount?: number | string
+}
+
+export class AutoApprovalHandler {
+	private consecutiveAutoApprovedRequestsCount: number = 0
+	private consecutiveAutoApprovedCost: number = 0
+
+	/**
+	 * Check if auto-approval limits have been reached and handle user approval if needed
+	 */
+	async checkAutoApprovalLimits(
+		state: GlobalState | undefined,
+		messages: ClineMessage[],
+		askForApproval: (
+			type: ClineAsk,
+			data: string,
+		) => Promise<{ response: ClineAskResponse; text?: string; images?: string[] }>,
+	): Promise<AutoApprovalResult> {
+		// Check request count limit
+		const requestResult = await this.checkRequestLimit(state, askForApproval)
+		if (!requestResult.shouldProceed || requestResult.requiresApproval) {
+			return requestResult
+		}
+
+		// Check cost limit
+		const costResult = await this.checkCostLimit(state, messages, askForApproval)
+		return costResult
+	}
+
+	/**
+	 * Increment the request counter and check if limit is exceeded
+	 */
+	private async checkRequestLimit(
+		state: GlobalState | undefined,
+		askForApproval: (
+			type: ClineAsk,
+			data: string,
+		) => Promise<{ response: ClineAskResponse; text?: string; images?: string[] }>,
+	): Promise<AutoApprovalResult> {
+		const maxRequests = state?.allowedMaxRequests || Infinity
+
+		// Increment the counter for each new API request
+		this.consecutiveAutoApprovedRequestsCount++
+
+		if (this.consecutiveAutoApprovedRequestsCount > maxRequests) {
+			const { response } = await askForApproval(
+				"auto_approval_max_req_reached",
+				JSON.stringify({ count: maxRequests, type: "requests" }),
+			)
+
+			// If we get past the promise, it means the user approved and did not start a new task
+			if (response === "yesButtonClicked") {
+				this.consecutiveAutoApprovedRequestsCount = 0
+				return {
+					shouldProceed: true,
+					requiresApproval: true,
+					approvalType: "requests",
+					approvalCount: maxRequests,
+				}
+			}
+
+			return {
+				shouldProceed: false,
+				requiresApproval: true,
+				approvalType: "requests",
+				approvalCount: maxRequests,
+			}
+		}
+
+		return { shouldProceed: true, requiresApproval: false }
+	}
+
+	/**
+	 * Calculate current cost and check if limit is exceeded
+	 */
+	private async checkCostLimit(
+		state: GlobalState | undefined,
+		messages: ClineMessage[],
+		askForApproval: (
+			type: ClineAsk,
+			data: string,
+		) => Promise<{ response: ClineAskResponse; text?: string; images?: string[] }>,
+	): Promise<AutoApprovalResult> {
+		const maxCost = state?.allowedMaxCost || Infinity
+
+		// Calculate total cost from messages
+		this.consecutiveAutoApprovedCost = getApiMetrics(messages).totalCost
+
+		// Use epsilon for floating-point comparison to avoid precision issues
+		const EPSILON = 0.0001
+		if (this.consecutiveAutoApprovedCost > maxCost + EPSILON) {
+			const { response } = await askForApproval(
+				"auto_approval_max_req_reached",
+				JSON.stringify({ count: maxCost.toFixed(2), type: "cost" }),
+			)
+
+			// If we get past the promise, it means the user approved and did not start a new task
+			if (response === "yesButtonClicked") {
+				// Note: We don't reset the cost to 0 here because the actual cost
+				// is calculated from the messages. This is different from the request count.
+				return {
+					shouldProceed: true,
+					requiresApproval: true,
+					approvalType: "cost",
+					approvalCount: maxCost.toFixed(2),
+				}
+			}
+
+			return {
+				shouldProceed: false,
+				requiresApproval: true,
+				approvalType: "cost",
+				approvalCount: maxCost.toFixed(2),
+			}
+		}
+
+		return { shouldProceed: true, requiresApproval: false }
+	}
+
+	/**
+	 * Reset the request counter (typically called when starting a new task)
+	 */
+	resetRequestCount(): void {
+		this.consecutiveAutoApprovedRequestsCount = 0
+	}
+
+	/**
+	 * Get current approval state for debugging/testing
+	 */
+	getApprovalState(): { requestCount: number; currentCost: number } {
+		return {
+			requestCount: this.consecutiveAutoApprovedRequestsCount,
+			currentCost: this.consecutiveAutoApprovedCost,
+		}
+	}
+}

src/core/task/Task.ts

@@ -92,6 +92,7 @@ import { ApiMessage } from "../task-persistence/apiMessages"
 import { getMessagesSinceLastSummary, summarizeConversation } from "../condense"
 import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
 import { restoreTodoListForTask } from "../tools/updateTodoListTool"
+import { AutoApprovalHandler } from "./AutoApprovalHandler"
 
 const MAX_EXPONENTIAL_BACKOFF_SECONDS = 600 // 10 minutes
 
@@ -199,7 +200,7 @@ export class Task extends EventEmitter<TaskEvents> {
 	readonly apiConfiguration: ProviderSettings
 	api: ApiHandler
 	private static lastGlobalApiRequestTime?: number
-	private consecutiveAutoApprovedRequestsCount: number = 0
+	private autoApprovalHandler: AutoApprovalHandler
 
 	/**
 	 * Reset the global API request timestamp. This should only be used for testing.
@@ -302,6 +303,7 @@ export class Task extends EventEmitter<TaskEvents> {
 
 		this.apiConfiguration = apiConfiguration
 		this.api = buildApiHandler(apiConfiguration)
+		this.autoApprovalHandler = new AutoApprovalHandler()
 
 		this.urlContentFetcher = new UrlContentFetcher(provider.context)
 		this.browserSession = new BrowserSession(provider.context)
@@ -1968,18 +1970,16 @@ export class Task extends EventEmitter<TaskEvents> {
 			({ role, content }) => ({ role, content }),
 		)
 
-		// Check if we've reached the maximum number of auto-approved requests
-		const maxRequests = state?.allowedMaxRequests || Infinity
-
-		// Increment the counter for each new API request
-		this.consecutiveAutoApprovedRequestsCount++
+		// Check auto-approval limits
+		const approvalResult = await this.autoApprovalHandler.checkAutoApprovalLimits(
+			state,
+			this.combineMessages(this.clineMessages.slice(1)),
+			async (type, data) => this.ask(type, data),
+		)
 
-		if (this.consecutiveAutoApprovedRequestsCount > maxRequests) {
-			const { response } = await this.ask("auto_approval_max_req_reached", JSON.stringify({ count: maxRequests }))
-			// If we get past the promise, it means the user approved and did not start a new task
-			if (response === "yesButtonClicked") {
-				this.consecutiveAutoApprovedRequestsCount = 0
-			}
+		if (!approvalResult.shouldProceed) {
+			// User did not approve, task should be aborted
+			throw new Error("Auto-approval limit reached and user did not approve continuation")
 		}
 
 		const metadata: ApiHandlerCreateMessageMetadata = {