From ed1e6db4b3d61814dcc12965ce55a00598291ba6 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Sun, 10 Aug 2025 00:20:35 +0000 Subject: [PATCH 1/2] fix: properly reset cost limit tracking when user clicks "Reset and Continue" - Track the message index at the time of cost reset - Only calculate costs from messages after the reset point - Prevents the issue where cost limit immediately triggers again after reset Fixes #6889 --- src/core/task/AutoApprovalHandler.ts | 14 +++-- .../__tests__/AutoApprovalHandler.spec.ts | 58 +++++++++++++++++-- 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/src/core/task/AutoApprovalHandler.ts b/src/core/task/AutoApprovalHandler.ts index 33821ddfa24..3054cd8b8d1 100644 --- a/src/core/task/AutoApprovalHandler.ts +++ b/src/core/task/AutoApprovalHandler.ts @@ -12,6 +12,7 @@ export interface AutoApprovalResult { export class AutoApprovalHandler { private consecutiveAutoApprovedRequestsCount: number = 0 private consecutiveAutoApprovedCost: number = 0 + private costResetMessageIndex: number = 0 /** * Check if auto-approval limits have been reached and handle user approval if needed @@ -91,8 +92,9 @@ export class AutoApprovalHandler { ): Promise { const maxCost = state?.allowedMaxCost || Infinity - // Calculate total cost from messages - this.consecutiveAutoApprovedCost = getApiMetrics(messages).totalCost + // Calculate total cost from messages after the last reset point + const messagesAfterReset = messages.slice(this.costResetMessageIndex) + this.consecutiveAutoApprovedCost = getApiMetrics(messagesAfterReset).totalCost // Use epsilon for floating-point comparison to avoid precision issues const EPSILON = 0.0001 @@ -104,8 +106,9 @@ export class AutoApprovalHandler { // 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. + // Reset the cost tracking by recording the current message count + // Future cost calculations will only include messages after this point + this.costResetMessageIndex = messages.length return { shouldProceed: true, requiresApproval: true, @@ -126,10 +129,11 @@ export class AutoApprovalHandler { } /** - * Reset the request counter (typically called when starting a new task) + * Reset the request counter and cost tracking (typically called when starting a new task) */ resetRequestCount(): void { this.consecutiveAutoApprovedRequestsCount = 0 + this.costResetMessageIndex = 0 } /** diff --git a/src/core/task/__tests__/AutoApprovalHandler.spec.ts b/src/core/task/__tests__/AutoApprovalHandler.spec.ts index e200948a33d..46de0fdb0e4 100644 --- a/src/core/task/__tests__/AutoApprovalHandler.spec.ts +++ b/src/core/task/__tests__/AutoApprovalHandler.spec.ts @@ -183,17 +183,67 @@ describe("AutoApprovalHandler", () => { expect(result3.requiresApproval).toBe(true) }) - it("should not reset cost to zero on approval", async () => { + it("should reset cost tracking on approval", async () => { + const messages: ClineMessage[] = [ + { type: "say", say: "api_req_started", text: '{"cost": 3.0}', ts: 1000 }, + { type: "say", say: "api_req_started", text: '{"cost": 3.0}', ts: 2000 }, + ] + + // First check - cost exceeds limit (6.0 > 5.0) + mockGetApiMetrics.mockReturnValue({ totalCost: 6.0 }) + mockAskForApproval.mockResolvedValue({ response: "yesButtonClicked" }) + + const result1 = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + expect(result1.shouldProceed).toBe(true) + expect(result1.requiresApproval).toBe(true) + + // Add more messages after reset + messages.push( + { type: "say", say: "api_req_started", text: '{"cost": 2.0}', ts: 3000 }, + { type: "say", say: "api_req_started", text: '{"cost": 1.0}', ts: 4000 }, + ) + + // Second check - should only count messages after reset (3.0 < 5.0) + mockGetApiMetrics.mockReturnValue({ totalCost: 3.0 }) + const result2 = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + + // Should not require approval since cost after reset is under limit + expect(result2.shouldProceed).toBe(true) + expect(result2.requiresApproval).toBe(false) + + // Verify it's only calculating cost from messages after reset point + expect(mockGetApiMetrics).toHaveBeenLastCalledWith(messages.slice(2)) + }) + + it("should track multiple cost resets correctly", async () => { const messages: ClineMessage[] = [] + // First cost limit hit + messages.push({ type: "say", say: "api_req_started", text: '{"cost": 6.0}', ts: 1000 }) mockGetApiMetrics.mockReturnValue({ totalCost: 6.0 }) mockAskForApproval.mockResolvedValue({ response: "yesButtonClicked" }) await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) - // Cost should still be calculated from messages, not reset - const state = handler.getApprovalState() - expect(state.currentCost).toBe(6.0) + // Add more messages + messages.push( + { type: "say", say: "api_req_started", text: '{"cost": 3.0}', ts: 2000 }, + { type: "say", say: "api_req_started", text: '{"cost": 3.0}', ts: 3000 }, + ) + + // Second cost limit hit (only counting from index 1) + mockGetApiMetrics.mockReturnValue({ totalCost: 6.0 }) + await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + + // Add more messages after second reset + messages.push({ type: "say", say: "api_req_started", text: '{"cost": 2.0}', ts: 4000 }) + + // Third check - should only count from last reset + mockGetApiMetrics.mockReturnValue({ totalCost: 2.0 }) + const result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + + expect(result.requiresApproval).toBe(false) + expect(mockGetApiMetrics).toHaveBeenLastCalledWith(messages.slice(3)) }) }) From 07283bfb9fd9a1fe497f2685400be1f1efa7f3e2 Mon Sep 17 00:00:00 2001 From: daniel-lxs Date: Thu, 14 Aug 2025 18:49:52 -0500 Subject: [PATCH 2/2] refactor: unify request and cost tracking approach in AutoApprovalHandler - Use consistent message-based tracking for both request count and cost limit - Track reset point using lastResetMessageIndex for both limits - Calculate request count from API messages similar to cost calculation - Simplify logic by removing separate tracking mechanisms - Update tests to match new message-based counting approach This addresses review feedback about inconsistent tracking approaches between request counting and cost tracking. --- src/core/task/AutoApprovalHandler.ts | 30 +++-- .../__tests__/AutoApprovalHandler.spec.ts | 103 ++++++++++++------ 2 files changed, 88 insertions(+), 45 deletions(-) diff --git a/src/core/task/AutoApprovalHandler.ts b/src/core/task/AutoApprovalHandler.ts index 3054cd8b8d1..db6ff04fe19 100644 --- a/src/core/task/AutoApprovalHandler.ts +++ b/src/core/task/AutoApprovalHandler.ts @@ -10,9 +10,9 @@ export interface AutoApprovalResult { } export class AutoApprovalHandler { + private lastResetMessageIndex: number = 0 private consecutiveAutoApprovedRequestsCount: number = 0 private consecutiveAutoApprovedCost: number = 0 - private costResetMessageIndex: number = 0 /** * Check if auto-approval limits have been reached and handle user approval if needed @@ -26,7 +26,7 @@ export class AutoApprovalHandler { ) => Promise<{ response: ClineAskResponse; text?: string; images?: string[] }>, ): Promise { // Check request count limit - const requestResult = await this.checkRequestLimit(state, askForApproval) + const requestResult = await this.checkRequestLimit(state, messages, askForApproval) if (!requestResult.shouldProceed || requestResult.requiresApproval) { return requestResult } @@ -37,10 +37,11 @@ export class AutoApprovalHandler { } /** - * Increment the request counter and check if limit is exceeded + * Calculate request count and check if limit is exceeded */ private async checkRequestLimit( state: GlobalState | undefined, + messages: ClineMessage[], askForApproval: ( type: ClineAsk, data: string, @@ -48,8 +49,11 @@ export class AutoApprovalHandler { ): Promise { const maxRequests = state?.allowedMaxRequests || Infinity - // Increment the counter for each new API request - this.consecutiveAutoApprovedRequestsCount++ + // Calculate request count from messages after the last reset point + const messagesAfterReset = messages.slice(this.lastResetMessageIndex) + // Count API request messages (simplified - you may need to adjust based on your message structure) + this.consecutiveAutoApprovedRequestsCount = + messagesAfterReset.filter((msg) => msg.type === "say" && msg.say === "api_req_started").length + 1 // +1 for the current request being checked if (this.consecutiveAutoApprovedRequestsCount > maxRequests) { const { response } = await askForApproval( @@ -59,7 +63,8 @@ export class AutoApprovalHandler { // If we get past the promise, it means the user approved and did not start a new task if (response === "yesButtonClicked") { - this.consecutiveAutoApprovedRequestsCount = 0 + // Reset tracking by recording the current message count + this.lastResetMessageIndex = messages.length return { shouldProceed: true, requiresApproval: true, @@ -93,7 +98,7 @@ export class AutoApprovalHandler { const maxCost = state?.allowedMaxCost || Infinity // Calculate total cost from messages after the last reset point - const messagesAfterReset = messages.slice(this.costResetMessageIndex) + const messagesAfterReset = messages.slice(this.lastResetMessageIndex) this.consecutiveAutoApprovedCost = getApiMetrics(messagesAfterReset).totalCost // Use epsilon for floating-point comparison to avoid precision issues @@ -106,9 +111,9 @@ export class AutoApprovalHandler { // If we get past the promise, it means the user approved and did not start a new task if (response === "yesButtonClicked") { - // Reset the cost tracking by recording the current message count - // Future cost calculations will only include messages after this point - this.costResetMessageIndex = messages.length + // Reset tracking by recording the current message count + // Future calculations will only include messages after this point + this.lastResetMessageIndex = messages.length return { shouldProceed: true, requiresApproval: true, @@ -129,11 +134,12 @@ export class AutoApprovalHandler { } /** - * Reset the request counter and cost tracking (typically called when starting a new task) + * Reset the tracking (typically called when starting a new task) */ resetRequestCount(): void { + this.lastResetMessageIndex = 0 this.consecutiveAutoApprovedRequestsCount = 0 - this.costResetMessageIndex = 0 + this.consecutiveAutoApprovedCost = 0 } /** diff --git a/src/core/task/__tests__/AutoApprovalHandler.spec.ts b/src/core/task/__tests__/AutoApprovalHandler.spec.ts index 46de0fdb0e4..eb9834ffeb2 100644 --- a/src/core/task/__tests__/AutoApprovalHandler.spec.ts +++ b/src/core/task/__tests__/AutoApprovalHandler.spec.ts @@ -40,12 +40,15 @@ describe("AutoApprovalHandler", () => { mockState.allowedMaxCost = 10 const messages: ClineMessage[] = [] - // First call should be under limit + // First call should be under limit (count = 1) const result1 = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) expect(result1.shouldProceed).toBe(true) expect(result1.requiresApproval).toBe(false) - // Second call should trigger request limit + // Add a message to simulate first request completed + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 1000 }) + + // Second call should trigger request limit (1 message + current = 2 > 1) mockAskForApproval.mockResolvedValue({ response: "yesButtonClicked" }) const result2 = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) @@ -64,27 +67,35 @@ describe("AutoApprovalHandler", () => { mockState.allowedMaxRequests = 3 }) - it("should increment request count on each check", async () => { + it("should calculate request count from messages", async () => { const messages: ClineMessage[] = [] - // Check state after each call - for (let i = 1; i <= 3; i++) { - await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) - const state = handler.getApprovalState() - expect(state.requestCount).toBe(i) - } + // First check - no messages yet, count should be 1 (for current request) + await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + let state = handler.getApprovalState() + expect(state.requestCount).toBe(1) + + // Add API request messages + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 1000 }) + await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + state = handler.getApprovalState() + expect(state.requestCount).toBe(2) // 1 message + current request + + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 2000 }) + await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + state = handler.getApprovalState() + expect(state.requestCount).toBe(3) // 2 messages + current request }) it("should ask for approval when limit is exceeded", async () => { const messages: ClineMessage[] = [] - // Make 3 requests (within limit) + // Add 3 API request messages (to simulate 3 requests made) for (let i = 0; i < 3; i++) { - await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 1000 + i }) } - expect(mockAskForApproval).not.toHaveBeenCalled() - // 4th request should trigger approval + // Next check should trigger approval (3 messages + current = 4 > 3) mockAskForApproval.mockResolvedValue({ response: "yesButtonClicked" }) const result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) @@ -99,29 +110,35 @@ describe("AutoApprovalHandler", () => { it("should reset count when user approves", async () => { const messages: ClineMessage[] = [] - // Exceed limit + // Add messages to exceed limit for (let i = 0; i < 3; i++) { - await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 1000 + i }) } - // 4th request should trigger approval and reset + // Next request should trigger approval and reset mockAskForApproval.mockResolvedValue({ response: "yesButtonClicked" }) await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) - // Count should be reset + // Add more messages after reset + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 4000 }) + + // Next check should only count messages after reset + const result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + expect(result.requiresApproval).toBe(false) // Should not require approval (1 message + current = 2 <= 3) + const state = handler.getApprovalState() - expect(state.requestCount).toBe(0) + expect(state.requestCount).toBe(2) // 1 message after reset + current request }) it("should not proceed when user rejects", async () => { const messages: ClineMessage[] = [] - // Exceed limit + // Add messages to exceed limit for (let i = 0; i < 3; i++) { - await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 1000 + i }) } - // 4th request with rejection + // Next request with rejection mockAskForApproval.mockResolvedValue({ response: "noButtonClicked" }) const result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) @@ -255,16 +272,21 @@ describe("AutoApprovalHandler", () => { mockGetApiMetrics.mockReturnValue({ totalCost: 3.0 }) - // First two requests should pass - for (let i = 0; i < 2; i++) { - const result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) - expect(result.shouldProceed).toBe(true) - expect(result.requiresApproval).toBe(false) - } + // First request should pass (count = 1) + let result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + expect(result.shouldProceed).toBe(true) + expect(result.requiresApproval).toBe(false) + + // Add a message and check again (count = 2) + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 1000 }) + result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + expect(result.shouldProceed).toBe(true) + expect(result.requiresApproval).toBe(false) - // Third request should trigger request limit (not cost limit) + // Add another message - third request should trigger request limit (count = 3 > 2) + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 2000 }) mockAskForApproval.mockResolvedValue({ response: "yesButtonClicked" }) - const result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + result = await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) expect(mockAskForApproval).toHaveBeenCalledWith( "auto_approval_max_req_reached", @@ -277,23 +299,38 @@ describe("AutoApprovalHandler", () => { }) describe("resetRequestCount", () => { - it("should reset the request counter", async () => { + it("should reset tracking", async () => { mockState.allowedMaxRequests = 5 + mockState.allowedMaxCost = 10.0 const messages: ClineMessage[] = [] - // Make some requests + // Add some messages for (let i = 0; i < 3; i++) { - await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + messages.push({ type: "say", say: "api_req_started", text: "{}", ts: 1000 + i }) } + mockGetApiMetrics.mockReturnValue({ totalCost: 5.0 }) + await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + let state = handler.getApprovalState() - expect(state.requestCount).toBe(3) + expect(state.requestCount).toBe(4) // 3 messages + current + expect(state.currentCost).toBe(5.0) // Reset handler.resetRequestCount() + // After reset, counts should be zero state = handler.getApprovalState() expect(state.requestCount).toBe(0) + expect(state.currentCost).toBe(0) + + // Next check should start fresh + mockGetApiMetrics.mockReturnValue({ totalCost: 8.0 }) + await handler.checkAutoApprovalLimits(mockState, messages, mockAskForApproval) + + state = handler.getApprovalState() + expect(state.requestCount).toBe(4) // All messages counted again + expect(state.currentCost).toBe(8.0) }) }) })