From 88cbaa304d3e896fdfb8ac30fbb2d19ee8b3705c Mon Sep 17 00:00:00 2001 From: Matt Bierner Date: Tue, 19 Aug 2025 13:40:15 -0700 Subject: [PATCH] Move getting auth token into `GithubCodeSearchService` (#673) This hides the implementation a bit better and should let us find the correct token for different repos --- .../extension/tools/node/githubRepoTool.tsx | 44 +++++--------- .../common/githubCodeSearchService.ts | 58 +++++++++++++------ .../common/remoteCodeSearch.ts | 5 ++ .../node/codeSearchRepoTracker.ts | 31 +++++----- .../node/codeSearchChunkSearch.ts | 2 +- .../base/simuliationWorkspaceChunkSearch.ts | 10 ++-- 6 files changed, 82 insertions(+), 68 deletions(-) diff --git a/extensions/copilot/src/extension/tools/node/githubRepoTool.tsx b/extensions/copilot/src/extension/tools/node/githubRepoTool.tsx index a3d5650eb82..1434a44b489 100644 --- a/extensions/copilot/src/extension/tools/node/githubRepoTool.tsx +++ b/extensions/copilot/src/extension/tools/node/githubRepoTool.tsx @@ -6,7 +6,6 @@ import * as l10n from '@vscode/l10n'; import { BasePromptElementProps, PromptElement, PromptElementProps, PromptPiece, PromptReference, PromptSizing } from '@vscode/prompt-tsx'; import type * as vscode from 'vscode'; -import { IAuthenticationService } from '../../../platform/authentication/common/authentication'; import { FileChunkAndScore } from '../../../platform/chunking/common/chunk'; import { IRunCommandExecutionService } from '../../../platform/commands/common/runCommandExecutionService'; import { GithubRepoId, toGithubNwo } from '../../../platform/git/common/gitService'; @@ -48,11 +47,9 @@ export class GithubRepoTool implements ICopilotTool { constructor( @IRunCommandExecutionService _commandService: IRunCommandExecutionService, @IInstantiationService private readonly _instantiationService: IInstantiationService, - @IAuthenticationService private readonly _authenticationService: IAuthenticationService, @IGithubCodeSearchService private readonly _githubCodeSearch: IGithubCodeSearchService, @ITelemetryService private readonly _telemetryService: ITelemetryService, - ) { - } + ) { } async invoke(options: vscode.LanguageModelToolInvocationOptions, token: CancellationToken): Promise { const githubRepoId = GithubRepoId.parse(options.input.repo); @@ -60,17 +57,12 @@ export class GithubRepoTool implements ICopilotTool { throw new Error('Invalid input. Could not parse repo'); } - const authToken = await this.tryGetAuthToken(); - if (!authToken) { - throw new Error('Not authenticated'); - } - const embeddingType = await this._availableEmbeddingTypesManager.value.getPreferredType(false); if (!embeddingType) { throw new Error('No embedding models available'); } - const searchResults = await this._githubCodeSearch.searchRepo(authToken, embeddingType, { githubRepoId, localRepoRoot: undefined, indexedCommit: undefined }, options.input.query, 64, {}, new TelemetryCorrelationId('github-repo-tool'), token); + const searchResults = await this._githubCodeSearch.searchRepo({ silent: true }, embeddingType, { githubRepoId, localRepoRoot: undefined, indexedCommit: undefined }, options.input.query, 64, {}, new TelemetryCorrelationId('github-repo-tool'), token); // Map the chunks to URIs // TODO: Won't work for proxima or branches not called main @@ -160,21 +152,20 @@ export class GithubRepoTool implements ICopilotTool { }); } - const authToken = await raceCancellationError(this.tryGetAuthToken(), token); - if (!authToken) { - return Result.error({ - message: l10n.t`Not authenticated`, - id: 'no-auth-token', - }); - } - const checkIndexReady = async (): Promise> => { - const state = await raceCancellationError(this._githubCodeSearch.getRemoteIndexState(authToken, githubRepoId, token), token); + const state = await raceCancellationError(this._githubCodeSearch.getRemoteIndexState({ silent: true }, githubRepoId, token), token); if (!state.isOk()) { - return Result.error({ - message: l10n.t`Could not check status of Github repo index`, - id: 'could-not-check-status', - }); + if (state.err.type === 'not-authorized') { + return Result.error({ + message: l10n.t`Not authenticated`, + id: 'no-auth-token', + }); + } else { + return Result.error({ + message: l10n.t`Could not check status of Github repo index`, + id: 'could-not-check-status', + }); + } } if (state.val.status === RemoteCodeSearchIndexStatus.Ready) { @@ -193,7 +184,7 @@ export class GithubRepoTool implements ICopilotTool { return Result.ok(githubRepoId); } - if (!await this._githubCodeSearch.triggerIndexing(authToken, 'tool', githubRepoId, new TelemetryCorrelationId('GitHubRepoTool'))) { + if (!await this._githubCodeSearch.triggerIndexing({ silent: true }, 'tool', githubRepoId, new TelemetryCorrelationId('GitHubRepoTool'))) { return Result.error({ message: l10n.t`Could not index Github repo. Repo may not exist or you may not have access to it.`, id: 'trigger-indexing-failed', @@ -214,11 +205,6 @@ export class GithubRepoTool implements ICopilotTool { id: 'not-ready-after-polling', }); } - - private async tryGetAuthToken() { - return (await this._authenticationService.getPermissiveGitHubSession({ silent: true }))?.accessToken - ?? (await this._authenticationService.getAnyGitHubSession({ silent: true }))?.accessToken; - } } interface GithubChunkSearchResultsProps extends BasePromptElementProps { diff --git a/extensions/copilot/src/platform/remoteCodeSearch/common/githubCodeSearchService.ts b/extensions/copilot/src/platform/remoteCodeSearch/common/githubCodeSearchService.ts index 369aaa02737..f1ab5a2b981 100644 --- a/extensions/copilot/src/platform/remoteCodeSearch/common/githubCodeSearchService.ts +++ b/extensions/copilot/src/platform/remoteCodeSearch/common/githubCodeSearchService.ts @@ -14,6 +14,7 @@ import { env } from '../../../util/vs/base/common/process'; import { URI } from '../../../util/vs/base/common/uri'; import { Range } from '../../../util/vs/editor/common/core/range'; import { createDecorator } from '../../../util/vs/platform/instantiation/common/instantiation'; +import { IAuthenticationService } from '../../authentication/common/authentication'; import { FileChunkAndScore } from '../../chunking/common/chunk'; import { getGithubMetadataHeaders } from '../../chunking/common/chunkingEndpointClientImpl'; import { stripChunkTextMetadata, truncateToMaxUtf8Length } from '../../chunking/common/chunkingStringUtils'; @@ -27,7 +28,7 @@ import { ILogService } from '../../log/common/logService'; import { IFetcherService, Response } from '../../networking/common/fetcherService'; import { postRequest } from '../../networking/common/networking'; import { ITelemetryService } from '../../telemetry/common/telemetry'; -import { CodeSearchOptions, CodeSearchResult, RemoteCodeSearchIndexState, RemoteCodeSearchIndexStatus } from './remoteCodeSearch'; +import { CodeSearchOptions, CodeSearchResult, RemoteCodeSearchError, RemoteCodeSearchIndexState, RemoteCodeSearchIndexStatus } from './remoteCodeSearch'; interface ResponseShape { @@ -70,16 +71,16 @@ export interface IGithubCodeSearchService { * Gets the state of the remote index for a given repo. */ getRemoteIndexState( - authToken: string, + authOptions: { readonly silent: boolean }, githubRepoId: GithubRepoId, token: CancellationToken, - ): Promise>; + ): Promise>; /** * Requests that a given repo be indexed. */ triggerIndexing( - authToken: string, + authOptions: { readonly silent: boolean }, triggerReason: 'auto' | 'manual' | 'tool', githubRepoId: GithubRepoId, telemetryInfo: TelemetryCorrelationId, @@ -91,7 +92,7 @@ export interface IGithubCodeSearchService { * The repo must have been indexed first. Make sure to check {@link getRemoteIndexState} or call {@link triggerIndexing}. */ searchRepo( - authToken: string, + authOptions: { readonly silent: boolean }, embeddingType: EmbeddingType, repo: GithubCodeSearchRepoInfo, query: string, @@ -107,6 +108,7 @@ export class GithubCodeSearchService implements IGithubCodeSearchService { declare readonly _serviceBrand: undefined; constructor( + @IAuthenticationService private readonly _authenticationService: IAuthenticationService, @IDomainService private readonly _domainService: IDomainService, @ICAPIClientService private readonly _capiClientService: ICAPIClientService, @IEnvService private readonly _envService: IEnvService, @@ -116,13 +118,19 @@ export class GithubCodeSearchService implements IGithubCodeSearchService { @ITelemetryService private readonly _telemetryService: ITelemetryService, ) { } - async getRemoteIndexState(authToken: string, githubRepoId: GithubRepoId, token: CancellationToken): Promise> { + async getRemoteIndexState(auth: { readonly silent: boolean }, githubRepoId: GithubRepoId, token: CancellationToken): Promise> { const repoNwo = toGithubNwo(githubRepoId); if (repoNwo.startsWith('microsoft/simuluation-test-')) { return Result.ok({ status: RemoteCodeSearchIndexStatus.NotYetIndexed }); } + const authToken = await this.getGithubAccessToken(auth.silent); + if (!authToken) { + this._logService.error(`GithubCodeSearchService::getRemoteIndexState(${repoNwo}). Failed to fetch indexing status. No valid github auth token.`); + return Result.error({ type: 'not-authorized' }); + } + try { const statusRequest = await raceCancellationError(this._capiClientService.makeRequest({ method: 'GET', @@ -142,14 +150,14 @@ export class GithubCodeSearchService implements IGithubCodeSearchService { statusCode: statusRequest.status, }); - this._logService.error(`GithubCodeSearchService.getRemoteIndexState(${repoNwo}). Failed to fetch indexing status. Response: ${statusRequest.status}. ${await statusRequest.text()}`); - return Result.error(new Error(`Failed to fetch indexing status. Response: ${statusRequest.status}.`)); + this._logService.error(`GithubCodeSearchService::getRemoteIndexState(${repoNwo}). Failed to fetch indexing status. Response: ${statusRequest.status}. ${await statusRequest.text()}`); + return Result.error({ type: 'generic-error', error: new Error(`Failed to fetch indexing status. Response: ${statusRequest.status}.`) }); } const preCheckResult = await raceCancellationError(statusRequest.json(), token); if (preCheckResult.semantic_code_search_ok && preCheckResult.semantic_commit_sha) { const indexedCommit = preCheckResult.semantic_commit_sha; - this._logService.trace(`GithubCodeSearchService.getRemoteIndexState(${repoNwo}). Found indexed commit: ${indexedCommit}.`); + this._logService.trace(`GithubCodeSearchService::getRemoteIndexState(${repoNwo}). Found indexed commit: ${indexedCommit}.`); return Result.ok({ status: RemoteCodeSearchIndexStatus.Ready, indexedCommit, @@ -158,36 +166,41 @@ export class GithubCodeSearchService implements IGithubCodeSearchService { if (preCheckResult.semantic_indexing_enabled) { if (await raceCancellationError(this.isEmptyRepo(authToken, githubRepoId, token), token)) { - this._logService.trace(`GithubCodeSearchService.getRemoteIndexState(${repoNwo}). Semantic indexing enabled but repo is empty.`); + this._logService.trace(`GithubCodeSearchService::getRemoteIndexState(${repoNwo}). Semantic indexing enabled but repo is empty.`); return Result.ok({ status: RemoteCodeSearchIndexStatus.Ready, indexedCommit: undefined }); } - this._logService.trace(`GithubCodeSearchService.getRemoteIndexState(${repoNwo}). Semantic indexing enabled but not yet indexed.`); + this._logService.trace(`GithubCodeSearchService::getRemoteIndexState(${repoNwo}). Semantic indexing enabled but not yet indexed.`); return Result.ok({ status: RemoteCodeSearchIndexStatus.BuildingIndex }); } else { - this._logService.trace(`GithubCodeSearchService.getRemoteIndexState(${repoNwo}). semantic_indexing_enabled was false. Repo not yet indexed but possibly can be.`); + this._logService.trace(`GithubCodeSearchService::getRemoteIndexState(${repoNwo}). semantic_indexing_enabled was false. Repo not yet indexed but possibly can be.`); return Result.ok({ status: RemoteCodeSearchIndexStatus.NotYetIndexed }); } - } catch (e) { + } catch (e: unknown) { if (isCancellationError(e)) { throw e; } - this._logService.error(`GithubCodeSearchService.getRemoteIndexState(${repoNwo}). Error: ${e}`); - return Result.error(e); + this._logService.error(`GithubCodeSearchService::getRemoteIndexState(${repoNwo}). Error: ${e}`); + return Result.error({ type: 'generic-error', error: e instanceof Error ? e : new Error(String(e)) }); } } public async triggerIndexing( - authToken: string, + auth: { readonly silent: boolean }, triggerReason: 'auto' | 'manual' | 'tool', githubRepoId: GithubRepoId, telemetryInfo: TelemetryCorrelationId, ): Promise { + const authToken = await this.getGithubAccessToken(auth.silent); + if (!authToken) { + return false; + } + const response = await this._capiClientService.makeRequest({ method: 'POST', headers: { @@ -241,7 +254,7 @@ export class GithubCodeSearchService implements IGithubCodeSearchService { } async searchRepo( - authToken: string, + auth: { readonly silent: boolean }, embeddingType: EmbeddingType, repo: GithubCodeSearchRepoInfo, searchQuery: string, @@ -250,6 +263,11 @@ export class GithubCodeSearchService implements IGithubCodeSearchService { telemetryInfo: TelemetryCorrelationId, token: CancellationToken ): Promise { + const authToken = await this.getGithubAccessToken(auth.silent); + if (!authToken) { + throw new Error('No valid auth token'); + } + const response = await raceCancellationError( postRequest( this._fetcherService, @@ -330,6 +348,12 @@ export class GithubCodeSearchService implements IGithubCodeSearchService { return result; } + private async getGithubAccessToken(silent: boolean) { + return (await this._authenticationService.getPermissiveGitHubSession({ silent }))?.accessToken + ?? (await this._authenticationService.getAnyGitHubSession({ silent }))?.accessToken; + } + + private async isEmptyRepo(authToken: string, githubRepoId: GithubRepoId, token: CancellationToken): Promise { const response = await raceCancellationError(fetch(this._capiClientService.dotcomAPIURL + `/repos/${toGithubNwo(githubRepoId)}`, { headers: { diff --git a/extensions/copilot/src/platform/remoteCodeSearch/common/remoteCodeSearch.ts b/extensions/copilot/src/platform/remoteCodeSearch/common/remoteCodeSearch.ts index 75ca5cebbcf..e53ba6da08f 100644 --- a/extensions/copilot/src/platform/remoteCodeSearch/common/remoteCodeSearch.ts +++ b/extensions/copilot/src/platform/remoteCodeSearch/common/remoteCodeSearch.ts @@ -22,7 +22,12 @@ export enum RemoteCodeSearchIndexStatus { export type RemoteCodeSearchIndexState = | { readonly status: RemoteCodeSearchIndexStatus.Ready; readonly indexedCommit: string | undefined } | { readonly status: RemoteCodeSearchIndexStatus.BuildingIndex | RemoteCodeSearchIndexStatus.NotYetIndexed | RemoteCodeSearchIndexStatus.NotIndexable } + ; +export type RemoteCodeSearchError = + | { readonly type: 'not-authorized' } + | { readonly type: 'generic-error'; readonly error: Error } + ; export interface CodeSearchResult { readonly chunks: readonly FileChunkAndScore[]; diff --git a/extensions/copilot/src/platform/remoteCodeSearch/node/codeSearchRepoTracker.ts b/extensions/copilot/src/platform/remoteCodeSearch/node/codeSearchRepoTracker.ts index 0f9d9a9a7f6..9d2898c4444 100644 --- a/extensions/copilot/src/platform/remoteCodeSearch/node/codeSearchRepoTracker.ts +++ b/extensions/copilot/src/platform/remoteCodeSearch/node/codeSearchRepoTracker.ts @@ -657,17 +657,21 @@ export class CodeSearchRepoTracker extends Disposable { let statusResult: Result; if (remoteInfo.repoId.type === 'github') { - const authToken = await this.getGithubAccessToken(true); - if (!authToken) { - this._logService.error(`CodeSearchRepoTracker.getIndexedStatus(${remoteInfo.repoId}).Failed to fetch indexing status.No valid github auth token.`); - return { - status: RepoStatus.NotAuthorized, - repo, - remoteInfo, - }; + const githubResult = await this._githubCodeSearchService.getRemoteIndexState({ silent: true }, remoteInfo.repoId, token); + if (githubResult.isOk()) { + statusResult = githubResult; + } else { + if (githubResult.err.type === 'not-authorized') { + this._logService.error(`CodeSearchRepoTracker.getIndexedStatus(${remoteInfo.repoId}). Failed to fetch indexing status. No valid github auth token.`); + return { + status: RepoStatus.NotAuthorized, + repo, + remoteInfo, + }; + } else { + statusResult = Result.error(githubResult.err.error); + } } - - statusResult = await this._githubCodeSearchService.getRemoteIndexState(authToken, remoteInfo.repoId, token); } else if (remoteInfo.repoId.type === 'ado') { const authToken = (await this._authenticationService.getAdoAccessTokenBase64({ silent: true })); if (!authToken) { @@ -708,11 +712,6 @@ export class CodeSearchRepoTracker extends Disposable { } } - private async getGithubAccessToken(silent: boolean) { - return (await this._authenticationService.getPermissiveGitHubSession({ silent }))?.accessToken - ?? (await this._authenticationService.getAnyGitHubSession({ silent }))?.accessToken; - } - private closeRepo(repo: RepoContext) { this._logService.trace(`CodeSearchRepoTracker.closeRepo(${repo.rootUri})`); @@ -836,7 +835,7 @@ export class CodeSearchRepoTracker extends Disposable { } const triggerSuccess = repoEntry.remoteInfo.repoId instanceof GithubRepoId - ? await this._githubCodeSearchService.triggerIndexing(authToken, triggerReason, repoEntry.remoteInfo.repoId, telemetryInfo) + ? await this._githubCodeSearchService.triggerIndexing({ silent: true }, triggerReason, repoEntry.remoteInfo.repoId, telemetryInfo) : await this._adoCodeSearchService.triggerIndexing(authToken, triggerReason, repoEntry.remoteInfo.repoId, telemetryInfo); if (this._isDisposed) { diff --git a/extensions/copilot/src/platform/workspaceChunkSearch/node/codeSearchChunkSearch.ts b/extensions/copilot/src/platform/workspaceChunkSearch/node/codeSearchChunkSearch.ts index bfbec9aeb4c..fab30c63003 100644 --- a/extensions/copilot/src/platform/workspaceChunkSearch/node/codeSearchChunkSearch.ts +++ b/extensions/copilot/src/platform/workspaceChunkSearch/node/codeSearchChunkSearch.ts @@ -569,7 +569,7 @@ export class CodeSearchChunkSearch extends Disposable implements IWorkspaceChunk return; } - return this._githubCodeSearchService.searchRepo(authToken, this._embeddingType, { + return this._githubCodeSearchService.searchRepo({ silent: true }, this._embeddingType, { githubRepoId: repo.remoteInfo.repoId, localRepoRoot: repo.repo.rootUri, indexedCommit: repo.status === RepoStatus.Ready ? repo.indexedCommit : undefined, diff --git a/extensions/copilot/test/base/simuliationWorkspaceChunkSearch.ts b/extensions/copilot/test/base/simuliationWorkspaceChunkSearch.ts index ba57bf05713..ce9b809f724 100644 --- a/extensions/copilot/test/base/simuliationWorkspaceChunkSearch.ts +++ b/extensions/copilot/test/base/simuliationWorkspaceChunkSearch.ts @@ -8,7 +8,7 @@ import { GithubRepoId } from '../../src/platform/git/common/gitService'; import { IIgnoreService } from '../../src/platform/ignore/common/ignoreService'; import { ILogService } from '../../src/platform/log/common/logService'; import { GithubCodeSearchRepoInfo, IGithubCodeSearchService, parseGithubCodeSearchResponse } from '../../src/platform/remoteCodeSearch/common/githubCodeSearchService'; -import { CodeSearchResult, RemoteCodeSearchIndexState, RemoteCodeSearchIndexStatus } from '../../src/platform/remoteCodeSearch/common/remoteCodeSearch'; +import { CodeSearchResult, RemoteCodeSearchError, RemoteCodeSearchIndexState, RemoteCodeSearchIndexStatus } from '../../src/platform/remoteCodeSearch/common/remoteCodeSearch'; import { BuildIndexTriggerReason, TriggerIndexingError } from '../../src/platform/remoteCodeSearch/node/codeSearchRepoTracker'; import { StrategySearchSizing, WorkspaceChunkQuery, WorkspaceChunkSearchOptions } from '../../src/platform/workspaceChunkSearch/common/workspaceChunkSearch'; import { FullWorkspaceChunkSearch } from '../../src/platform/workspaceChunkSearch/node/fullWorkspaceChunkSearch'; @@ -35,7 +35,7 @@ class SimulationGithubCodeSearchService extends Disposable implements IGithubCod super(); } - async searchRepo(authToken: string, embeddingType: EmbeddingType, repo: GithubCodeSearchRepoInfo, query: string, maxResults: number, options: WorkspaceChunkSearchOptions, _telemetryInfo: TelemetryCorrelationId, token: CancellationToken): Promise { + async searchRepo(authOptions: { silent: boolean }, embeddingType: EmbeddingType, repo: GithubCodeSearchRepoInfo, query: string, maxResults: number, options: WorkspaceChunkSearchOptions, _telemetryInfo: TelemetryCorrelationId, token: CancellationToken): Promise { this._logService.trace(`SimulationGithubCodeSearchService::searchRepo(${repo.githubRepoId}, ${query})`); const response = await fetch(searchEndpoint, { method: 'POST', @@ -62,11 +62,11 @@ class SimulationGithubCodeSearchService extends Disposable implements IGithubCod return result; } - async getRemoteIndexState(authToken: string, githubRepoId: GithubRepoId, token: CancellationToken): Promise> { + async getRemoteIndexState(authOptions: { silent: boolean }, githubRepoId: GithubRepoId, token: CancellationToken): Promise> { return Result.ok({ status: RemoteCodeSearchIndexStatus.Ready, indexedCommit: 'HEAD' }); } - triggerIndexing(authToken: string, triggerReason: 'auto' | 'manual' | 'tool', githubRepoId: GithubRepoId): Promise { + triggerIndexing(authOptions: { silent: boolean }, triggerReason: 'auto' | 'manual' | 'tool', githubRepoId: GithubRepoId): Promise { throw new Error('Method not implemented.'); } @@ -113,7 +113,7 @@ export class SimulationCodeSearchChunkSearchService extends Disposable implement const repo = new GithubRepoId('test-org', 'test-repo'); try { - const results = await this._githubCodeSearchService.searchRepo('', EmbeddingType.text3small_512, { + const results = await this._githubCodeSearchService.searchRepo({ silent: true }, EmbeddingType.text3small_512, { githubRepoId: repo, indexedCommit: undefined, localRepoRoot: undefined,