Skip to content

Back off token recount selection retries #377

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
baron wants to merge 1 commit into
base: main
Choose a base branch
from
+89 −1
Open

Back off token recount selection retries #377

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 61 additions & 1 deletion Sources/RepoPrompt/Features/Prompt/ViewModels/TokenCountingViewModel.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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)?
Expand Down Expand Up @@ -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)
Comment on lines +286 to +293

@sentry sentry Bot Jul 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: The selection consistency retry backoff can be canceled by concurrent non-selection dirty events (e.g., prompt text changes), defeating the purpose of the backoff delay.
Severity: MEDIUM

Suggested Fix

The shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount flag should not be reset until the backoff-delayed task actually begins execution. Alternatively, when scheduleTokenCountUpdateIfNeeded() is called, it should check if the pending tokenUpdateDebounceTask is a backoff task and avoid canceling it for non-critical dirty events.

Prompt for AI Agent 
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.

Location:
Sources/RepoPrompt/Features/Prompt/ViewModels/TokenCountingViewModel.swift#L286-L293

Potential issue: A race condition exists where the selection consistency retry backoff
is defeated. When a retry is scheduled, the
`shouldUseSelectionConsistencyRetryBackoffForNextSelectionRecount` flag is consumed and
a long debounce task is started. If a concurrent non-selection event (e.g., prompt text
change) triggers `markDirty()` before the backoff task executes, the pending backoff
task is canceled and replaced with a standard, shorter debounce task. This undermines
the feature's goal of reducing rapid recounts during periods of instability, leading to
degraded performance.

Did we get this right? 👍 / 👎 to inform future reviews.

} catch {
return
}
Expand All @@ -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,
Expand Down Expand Up @@ -348,6 +381,7 @@ class TokenCountingViewModel: ObservableObject {

private func recordSelectionProjectionChanged() {
selectionObservationRevision &+= 1
resetSelectionConsistencyRetry()
markDirty(.selection)
}

Expand Down Expand Up @@ -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)",
Expand Down Expand Up @@ -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",
Expand Down
28 changes: 28 additions & 0 deletions Tests/RepoPromptTests/MCP/MCPSelectionReplyFreshnessTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()) }
Expand Down
Loading