diff --git a/Sources/RepoPrompt/Features/Prompt/ViewModels/TokenCountingViewModel.swift b/Sources/RepoPrompt/Features/Prompt/ViewModels/TokenCountingViewModel.swift index 5a8ac9305..db573b6a5 100644 --- a/Sources/RepoPrompt/Features/Prompt/ViewModels/TokenCountingViewModel.swift +++ b/Sources/RepoPrompt/Features/Prompt/ViewModels/TokenCountingViewModel.swift @@ -73,6 +73,8 @@ class TokenCountingViewModel: ObservableObject { // MARK: - Private Properties private static let tokenUpdateDebounceNanoseconds: UInt64 = 500_000_000 + private static let maxSelectionConsistencyRetryCount = 5 + private static let maxSelectionConsistencyRetryDelayNanoseconds: UInt64 = 5_000_000_000 private let tokenCalculationService = TokenCalculationService() private let promptContextAccountingService = PromptContextAccountingService() @@ -86,6 +88,8 @@ class TokenCountingViewModel: ObservableObject { private var lastObservedSelectionObservationRevision: UInt64 = 0 private var lastPredominantLanguage: String = "Swift" private var automaticRecountSuspendDepth: Int = 0 + private var selectionConsistencyRetryCount = 0 + private var shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount = false private var lastPublishedSelection: StoredSelection? #if DEBUG private var debugBeforeTokenCalculationForTesting: (@MainActor @Sendable () async -> Void)? @@ -277,10 +281,16 @@ class TokenCountingViewModel: ObservableObject { } let generation = tokenCountSchedulerGeneration + let usesSelectionConsistencyRetryBackoff = shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount + && pendingDirty.contains(.selection) + shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount = false + let debounceNanoseconds = tokenCountUpdateDebounceNanoseconds( + useSelectionConsistencyRetryBackoff: usesSelectionConsistencyRetryBackoff + ) tokenUpdateDebounceTask?.cancel() tokenUpdateDebounceTask = Task { @MainActor [weak self] in do { - try await Task.sleep(nanoseconds: Self.tokenUpdateDebounceNanoseconds) + try await Task.sleep(nanoseconds: debounceNanoseconds) } catch { return } @@ -295,6 +305,29 @@ class TokenCountingViewModel: ObservableObject { } } + private func tokenCountUpdateDebounceNanoseconds(useSelectionConsistencyRetryBackoff: Bool) -> UInt64 { + guard useSelectionConsistencyRetryBackoff, + selectionConsistencyRetryCount > 0 + else { + return Self.tokenUpdateDebounceNanoseconds + } + let retryDelay = UInt64(selectionConsistencyRetryCount) * Self.tokenUpdateDebounceNanoseconds * 2 + return min(retryDelay, Self.maxSelectionConsistencyRetryDelayNanoseconds) + } + + private func recordSelectionConsistencyRetry() { + selectionConsistencyRetryCount = min( + selectionConsistencyRetryCount + 1, + Self.maxSelectionConsistencyRetryCount + ) + shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount = true + } + + private func resetSelectionConsistencyRetry() { + selectionConsistencyRetryCount = 0 + shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount = false + } + private func startPendingTokenCountUpdate(generation: UInt64) { guard isTokenCountSchedulerActive, tokenCountSchedulerGeneration == generation, @@ -348,6 +381,7 @@ class TokenCountingViewModel: ObservableObject { private func recordSelectionProjectionChanged() { selectionObservationRevision &+= 1 + resetSelectionConsistencyRetry() markDirty(.selection) } @@ -424,6 +458,30 @@ class TokenCountingViewModel: ObservableObject { debugTokenCalculationStartCount } + func selectionConsistencyRetryCountForTesting() -> Int { + selectionConsistencyRetryCount + } + + func selectionConsistencyRetryBackoffPendingForTesting() -> Bool { + shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount + } + + func recordSelectionConsistencyRetryForTesting() { + recordSelectionConsistencyRetry() + } + + func recordSelectionProjectionChangedForTesting() { + recordSelectionProjectionChanged() + } + + func tokenCountUpdateDebounceNanosecondsForTesting( + useSelectionConsistencyRetryBackoff: Bool = true + ) -> UInt64 { + tokenCountUpdateDebounceNanoseconds( + useSelectionConsistencyRetryBackoff: useSelectionConsistencyRetryBackoff + ) + } + func debugTokenRecountStateFields() -> [String: String] { [ "pendingDirtyRaw": "\(pendingDirty.rawValue)", @@ -772,9 +830,11 @@ class TokenCountingViewModel: ObservableObject { fields: ["duration": consistencyStartMS.map { PromptTokenRecountDiagnostics.formatElapsedMS(since: $0) } ?? "notMeasured"] ) #endif + recordSelectionConsistencyRetry() markDirty(.selection) return } + resetSelectionConsistencyRetry() #if DEBUG PromptTokenRecountDiagnostics.event( "tokenRecount.calculate.selectionConsistent", diff --git a/Tests/RepoPromptTests/MCP/MCPSelectionReplyFreshnessTests.swift b/Tests/RepoPromptTests/MCP/MCPSelectionReplyFreshnessTests.swift index ecc5b23a9..c0a416b79 100644 --- a/Tests/RepoPromptTests/MCP/MCPSelectionReplyFreshnessTests.swift +++ b/Tests/RepoPromptTests/MCP/MCPSelectionReplyFreshnessTests.swift @@ -845,6 +845,34 @@ final class MCPSelectionReplyFreshnessTests: XCTestCase { } #if DEBUG + func testSelectionConsistencyRetryDebounceBacksOffOnlyForRetrySelectionRecounts() { + let tokenCounter = TokenCountingViewModel() + + XCTAssertEqual(tokenCounter.tokenCountUpdateDebounceNanosecondsForTesting(), 500_000_000) + + tokenCounter.recordSelectionConsistencyRetryForTesting() + XCTAssertEqual(tokenCounter.selectionConsistencyRetryCountForTesting(), 1) + XCTAssertTrue(tokenCounter.selectionConsistencyRetryBackoffPendingForTesting()) + XCTAssertEqual( + tokenCounter.tokenCountUpdateDebounceNanosecondsForTesting( + useSelectionConsistencyRetryBackoff: false + ), + 500_000_000 + ) + XCTAssertEqual(tokenCounter.tokenCountUpdateDebounceNanosecondsForTesting(), 1_000_000_000) + + for _ in 0 ..< 6 { + tokenCounter.recordSelectionConsistencyRetryForTesting() + } + XCTAssertEqual(tokenCounter.selectionConsistencyRetryCountForTesting(), 5) + XCTAssertEqual(tokenCounter.tokenCountUpdateDebounceNanosecondsForTesting(), 5_000_000_000) + + tokenCounter.recordSelectionProjectionChangedForTesting() + XCTAssertEqual(tokenCounter.selectionConsistencyRetryCountForTesting(), 0) + XCTAssertFalse(tokenCounter.selectionConsistencyRetryBackoffPendingForTesting()) + XCTAssertEqual(tokenCounter.tokenCountUpdateDebounceNanosecondsForTesting(), 500_000_000) + } + func testActiveMCPTokenRepliesCompleteWhileBackgroundRecountIsBlockedAndCoalesce() async throws { let root = try makeTemporaryRoot(name: "ActiveTokenCache") defer { try? FileManager.default.removeItem(at: root.deletingLastPathComponent()) }