Merge pull request #208155 from microsoft/uninterested-shrew

fix cancellations on ai search
microsoft · Mar 20, 2024 · 6a52a46 · 6a52a46
2 parents 7fca5d7 + fa464e8
commit 6a52a46
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/src/vs/workbench/contrib/search/browser/searchModel.ts b/src/vs/workbench/contrib/search/browser/searchModel.ts
@@ -2013,6 +2013,7 @@ export class SearchModel extends Disposable {
  private currentCancelTokenSource: CancellationTokenSource | null = null;
  private currentAICancelTokenSource: CancellationTokenSource | null = null;
  private searchCancelledForNewSearch: boolean = false;
+ private aiSearchCancelledForNewSearch: boolean = false;
  public location: SearchModelLocation = SearchModelLocation.PANEL;
 
  constructor(
@@ -2297,11 +2298,15 @@ export class SearchModel extends Disposable {
  private onSearchError(e: any, duration: number, ai: boolean): void {
  if (errors.isCancellationError(e)) {
  this.onSearchCompleted(
- this.searchCancelledForNewSearch
+ (ai ? this.aiSearchCancelledForNewSearch : this.searchCancelledForNewSearch)
  ? { exit: SearchCompletionExitCode.NewSearchStarted, results: [], messages: [] }
  : undefined,
  duration, '', ai);
- this.searchCancelledForNewSearch = false;
+ if (ai) {
+ this.aiSearchCancelledForNewSearch = false;
+ } else {
+ this.searchCancelledForNewSearch = false;
+ }
  }
  }
 
@@ -2338,9 +2343,17 @@ export class SearchModel extends Disposable {
  }
  return false;
  }
-
+ cancelAISearch(cancelledForNewSearch = false): boolean {
+ if (this.currentAICancelTokenSource) {
+ this.aiSearchCancelledForNewSearch = cancelledForNewSearch;
+ this.currentAICancelTokenSource.cancel();
+ return true;
+ }
+ return false;
+ }
  override dispose(): void {
  this.cancelSearch();
+ this.cancelAISearch();
  this.searchResult.dispose();
  super.dispose();
  }

diff --git a/src/vs/workbench/contrib/search/browser/searchView.ts b/src/vs/workbench/contrib/search/browser/searchView.ts
@@ -332,6 +332,8 @@ export class SearchView extends ViewPane {
  return;
  }
 
+ // in each case, we want to cancel our current AI search because it is no longer valid
+ this.model.cancelAISearch();
  if (visible) {
  await this.model.addAIResults();
  } else {
@@ -1634,6 +1636,7 @@ export class SearchView extends ViewPane {
  });
 
  this.viewModel.cancelSearch(true);
+ this.viewModel.cancelAISearch(true);
 
  this.currentSearchQ = this.currentSearchQ
  .then(() => this.doSearch(query, excludePatternText, includePatternText, triggeredOnType))