Skip to content
This repository was archived by the owner on May 20, 2026. It is now read-only.

Commit 6bed187

Browse files
mjbvzCopilot
mjbvz
and
Copilot
authored
Add telemetry headers to external ingest requests (#3709)
* Add telemetry headers to external ingest requests * Update src/platform/workspaceChunkSearch/node/codeSearch/externalIngestClient.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 10cc8a0 commit 6bed187

6 files changed

Lines changed: 73 additions & 52 deletions

File tree

‎src/platform/workspaceChunkSearch/node/codeSearch/codeSearchChunkSearch.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import * as l10n from '@vscode/l10n';
77
import { shouldInclude } from '../../../../util/common/glob';
88
import { Result } from '../../../../util/common/result';
9-
import { TelemetryCorrelationId } from '../../../../util/common/telemetryCorrelationId';
9+
import { CallTracker, TelemetryCorrelationId } from '../../../../util/common/telemetryCorrelationId';
1010
import { coalesce } from '../../../../util/vs/base/common/arrays';
1111
import { IntervalTimer, raceCancellationError, raceTimeout } from '../../../../util/vs/base/common/async';
1212
import { CancellationToken, CancellationTokenSource } from '../../../../util/vs/base/common/cancellation';
@@ -557,7 +557,7 @@ export class CodeSearchChunkSearch extends Disposable implements IWorkspaceChunk
557557

558558
// Also search external ingest for files not in code search repos (if enabled)
559559
const externalIngestOperation = this.isExternalIngestEnabled()
560-
? this._externalIngestIndex.value.search(sizing, query, token).catch(e => {
560+
? this._externalIngestIndex.value.search(sizing, query, innerTelemetryInfo.callTracker, token).catch(e => {
561561
if (!isCancellationError(e)) {
562562
this._logService.warn(`External ingest search failed: ${e}`);
563563
}
@@ -857,7 +857,7 @@ export class CodeSearchChunkSearch extends Disposable implements IWorkspaceChunk
857857
// Update external ingest index if enabled
858858
const externalIndexEnabled = this.isExternalIngestEnabled();
859859
if (externalIndexEnabled) {
860-
const result = await raceCancellationError(this._externalIngestIndex.value.doIngest(onProgress, token), token);
860+
const result = await raceCancellationError(this._externalIngestIndex.value.doIngest(telemetryInfo.callTracker, onProgress, token), token);
861861
if (result.isError()) {
862862
return Result.error(result.err);
863863
}
@@ -1006,7 +1006,7 @@ export class CodeSearchChunkSearch extends Disposable implements IWorkspaceChunk
10061006
}
10071007
}
10081008

1009-
public deleteExternalIngestWorkspaceIndex(token: CancellationToken): Promise<void> {
1010-
return this._externalIngestIndex.value.deleteIndex(token);
1009+
public deleteExternalIngestWorkspaceIndex(callTracker: CallTracker, token: CancellationToken): Promise<void> {
1010+
return this._externalIngestIndex.value.deleteIndex(callTracker, token);
10111011
}
10121012
}

‎src/platform/workspaceChunkSearch/node/codeSearch/externalIngestApi.ts‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,13 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6+
import { CallTracker } from '../../../../util/common/telemetryCorrelationId';
67
import { raceCancellationError } from '../../../../util/vs/base/common/async';
78
import { CancellationToken } from '../../../../util/vs/base/common/cancellation';
89
import { CancellationError } from '../../../../util/vs/base/common/errors';
910
import { IDisposable } from '../../../../util/vs/base/common/lifecycle';
11+
import { getGithubMetadataHeaders } from '../../../chunking/common/chunkingEndpointClientImpl';
12+
import { IEnvService } from '../../../env/common/envService';
1013
import { ILogService } from '../../../log/common/logService';
1114

1215
// Sliding window that holds at least N entries and all entries in the time window.
@@ -202,6 +205,7 @@ export class ApiClient implements IDisposable {
202205

203206
constructor(
204207
target: number | null = 80,
208+
@IEnvService private readonly envService: IEnvService,
205209
@ILogService private readonly logService: ILogService,
206210
) {
207211
if (target === null) {
@@ -216,6 +220,7 @@ export class ApiClient implements IDisposable {
216220
headers: Record<string, string>,
217221
method: string,
218222
body: unknown | undefined,
223+
callerInfo: CallTracker,
219224
token: CancellationToken,
220225
): Promise<Response> {
221226
if (this.throttler) {
@@ -234,7 +239,10 @@ export class ApiClient implements IDisposable {
234239
try {
235240
const res = await fetch(url, {
236241
method,
237-
headers,
242+
headers: {
243+
...headers,
244+
...getGithubMetadataHeaders(callerInfo, this.envService),
245+
},
238246
body: body ? JSON.stringify(body) : undefined,
239247
});
240248
if (!res.ok) {

‎src/platform/workspaceChunkSearch/node/codeSearch/externalIngestClient.ts‎

Lines changed: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import * as l10n from '@vscode/l10n';
88
import crypto from 'crypto';
99
import { CancellationToken } from 'vscode-languageserver-protocol';
1010
import { Result } from '../../../../util/common/result';
11+
import { CallTracker } from '../../../../util/common/telemetryCorrelationId';
1112
import { raceCancellationError } from '../../../../util/vs/base/common/async';
1213
import { encodeBase64, VSBuffer } from '../../../../util/vs/base/common/buffer';
1314
import { CancellationError } from '../../../../util/vs/base/common/errors';
@@ -40,14 +41,15 @@ export interface IExternalIngestClient {
4041
filesetName: string,
4142
currentCheckpoint: string | undefined,
4243
allFiles: AsyncIterable<ExternalIngestFile>,
44+
callTracker: CallTracker,
4345
token: CancellationToken,
4446
onProgress?: (message: string) => void
4547
): Promise<Result<{ checkpoint: string }, Error>>;
4648

47-
listFilesets(token: CancellationToken): Promise<string[]>;
48-
deleteFileset(filesetName: string, token: CancellationToken): Promise<void>;
49+
listFilesets(callTracker: CallTracker, token: CancellationToken): Promise<string[]>;
50+
deleteFileset(filesetName: string, callTracker: CallTracker, token: CancellationToken): Promise<void>;
4951

50-
searchFilesets(filesetName: string, rootUri: URI, prompt: string, limit: number, token: CancellationToken): Promise<CodeSearchResult>;
52+
searchFilesets(filesetName: string, rootUri: URI, prompt: string, limit: number, callTracker: CallTracker, token: CancellationToken): Promise<CodeSearchResult>;
5153

5254
/**
5355
* Quickly checks if a file can be ingested based on its path and size.
@@ -106,10 +108,10 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
106108
return headers;
107109
}
108110

109-
private async post(authToken: string, path: string, body: unknown, options: { retries?: number }, token: CancellationToken): Promise<Response> {
111+
private async post(authToken: string, path: string, body: unknown, options: { retries?: number }, callTracker: CallTracker, token: CancellationToken): Promise<Response> {
110112
const retries = options.retries ?? 0;
111113
const url = `${ExternalIngestClient.baseUrl}${path}`;
112-
const response = await this.apiClient.makeRequest(url, this.getHeaders(authToken), 'POST', body, token);
114+
const response = await this.apiClient.makeRequest(url, this.getHeaders(authToken), 'POST', body, callTracker, token);
113115

114116
// Retry on 500 errors as these are often transient
115117
const shouldRetry = response.status.toString().startsWith('5') && retries > 0;
@@ -127,7 +129,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
127129

128130
if (shouldRetry) {
129131
this.logService.warn(`ExternalIngestClient::post(${path}): Got ${response.status}, retrying... (${retries} retries remaining)`);
130-
return this.post(authToken, path, body, { retries: retries - 1 }, token);
132+
return this.post(authToken, path, body, { retries: retries - 1 }, callTracker, token);
131133
}
132134

133135
if (!response.ok) {
@@ -138,7 +140,8 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
138140
return response;
139141
}
140142

141-
async updateIndex(filesetName: string, currentCheckpoint: string | undefined, allFiles: AsyncIterable<ExternalIngestFile>, token: CancellationToken, onProgress?: (message: string) => void): Promise<Result<{ checkpoint: string }, Error>> {
143+
async updateIndex(filesetName: string, currentCheckpoint: string | undefined, allFiles: AsyncIterable<ExternalIngestFile>, inCallTracker: CallTracker, token: CancellationToken, onProgress?: (message: string) => void): Promise<Result<{ checkpoint: string }, Error>> {
144+
const callTracker = inCallTracker.add('ExternalIngestClient::updateIndex');
142145
const authToken = await raceCancellationError(this.getAuthToken(), token);
143146
if (!authToken) {
144147
this.logService.warn('ExternalIngestClient::updateIndex(): No auth token available');
@@ -196,7 +199,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
196199
new_checkpoint: newCheckpoint,
197200
geo_filter: Buffer.from(geoFilter.toBytes()).toString('base64'),
198201
coded_symbols: codedSymbols,
199-
}, {}, token);
202+
}, {}, callTracker, token);
200203
};
201204

202205
let createIngestResponse: Response;
@@ -211,7 +214,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
211214
this.logService.info('ExternalIngestClient::updateIndex(): Got 429, cleaning up old filesets...');
212215
onProgress?.(l10n.t("Too many filesets, cleaning up old ones..."));
213216

214-
await raceCancellationError(this.cleanupOldFilesets(authToken, filesetName, token), token);
217+
await raceCancellationError(this.cleanupOldFilesets(authToken, filesetName, callTracker, token), token);
215218

216219
// Retry the create ingest
217220
this.logService.info('ExternalIngestClient::updateIndex(): Retrying create ingest after cleanup...');
@@ -281,6 +284,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
281284
coded_symbol_range: codedSymbolRange,
282285
},
283286
{},
287+
callTracker,
284288
token
285289
);
286290
const body = await raceCancellationError(pushCodedSymbolsResponse.json(), token) as { next_coded_symbol_range?: CodedSymbolRange };
@@ -320,7 +324,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
320324
const getBatchResponse = await this.post(authToken, '/external/code/ingest/batch', {
321325
ingest_id: ingestId,
322326
page_token: pageToken,
323-
}, {}, token);
327+
}, {}, callTracker, token);
324328

325329
const { doc_ids: docIds, next_page_token: nextPageToken } =
326330
await raceCancellationError(getBatchResponse.json(), token) as { doc_ids: string[] | undefined; next_page_token: string | undefined };
@@ -358,7 +362,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
358362
content,
359363
file_path: fileEntry.relativePath,
360364
doc_id: requestedDocSha,
361-
}, { retries: 3 }, token);
365+
}, { retries: 3 }, callTracker, token);
362366
if (!res.ok) {
363367
const requestId = res.headers.get(githubHeaders.requestId);
364368
const responseBody = await res.text();
@@ -403,7 +407,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
403407
onProgress?.(l10n.t('Finalizing index...'));
404408
const resp = await this.post(authToken, '/external/code/ingest/finalize', {
405409
ingest_id: ingestId,
406-
}, {}, token);
410+
}, {}, callTracker, token);
407411

408412
this.logService.info('ExternalIngestClient::updateIndex(): Successfully finalized ingest.');
409413
const requestId = resp.headers.get('x-github-request-id');
@@ -413,23 +417,24 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
413417
return Result.ok({ checkpoint: newCheckpoint });
414418
}
415419

416-
async listFilesets(token: CancellationToken): Promise<string[]> {
420+
async listFilesets(callTracker: CallTracker, token: CancellationToken): Promise<string[]> {
417421
const authToken = await this.getAuthToken();
418422
if (!authToken) {
419423
this.logService.warn('ExternalIngestClient::listFilesets(): No auth token available');
420424
return [];
421425
}
422426

423-
const filesets = await this.listFilesetsWithDetails(authToken, token);
427+
const filesets = await this.listFilesetsWithDetails(authToken, callTracker.add('ExternalIngestClient::listFilesets'), token);
424428
return filesets.map(x => x.name);
425429
}
426430

427-
private async listFilesetsWithDetails(authToken: string, token: CancellationToken): Promise<Array<{ name: string; checkpoint: string; status: string }>> {
431+
private async listFilesetsWithDetails(authToken: string, callTracker: CallTracker, token: CancellationToken): Promise<Array<{ name: string; checkpoint: string; status: string }>> {
428432
const resp = await this.apiClient.makeRequest(
429433
`${ExternalIngestClient.baseUrl}/external/code/ingest`,
430434
this.getHeaders(authToken),
431435
'GET',
432436
undefined,
437+
callTracker.add('ExternalIngestClient::listFilesetsWithDetails'),
433438
token
434439
);
435440

@@ -440,34 +445,36 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
440445
/**
441446
* Cleans up old filesets to make room for new ones.
442447
*/
443-
private async cleanupOldFilesets(authToken: string, currentFilesetName: string, token: CancellationToken): Promise<void> {
444-
const filesets = await this.listFilesetsWithDetails(authToken, token);
448+
private async cleanupOldFilesets(authToken: string, currentFilesetName: string, inCallTracker: CallTracker, token: CancellationToken): Promise<void> {
449+
const callTracker = inCallTracker.add('ExternalIngestClient::cleanupOldFilesets');
450+
const filesets = await this.listFilesetsWithDetails(authToken, callTracker, token);
445451

446452
const candidates = filesets.filter(f => f.name !== currentFilesetName);
447453
const toDelete = candidates.at(-1);
448454
if (toDelete) {
449-
await this.deleteFilesetByName(authToken, toDelete.name, token);
455+
await this.deleteFilesetByName(authToken, toDelete.name, callTracker, token);
450456
}
451457
}
452458

453-
async deleteFileset(filesetName: string, token: CancellationToken): Promise<void> {
459+
async deleteFileset(filesetName: string, callTracker: CallTracker, token: CancellationToken): Promise<void> {
454460
const authToken = await this.getAuthToken();
455461
if (!authToken) {
456462
this.logService.warn('ExternalIngestClient::deleteFileset(): No auth token available');
457463
return;
458464
}
459465

460-
return this.deleteFilesetByName(authToken, filesetName, token);
466+
return this.deleteFilesetByName(authToken, filesetName, callTracker.add('ExternalIngestClient::deleteFileset'), token);
461467
}
462468

463-
async deleteFilesetByName(authToken: string, fileSetName: string, token: CancellationToken): Promise<void> {
469+
async deleteFilesetByName(authToken: string, fileSetName: string, callTracker: CallTracker, token: CancellationToken): Promise<void> {
464470
const resp = await this.apiClient.makeRequest(
465471
`${ExternalIngestClient.baseUrl}/external/code/ingest`,
466472
this.getHeaders(authToken),
467473
'DELETE',
468474
{
469475
fileset_name: fileSetName,
470476
},
477+
callTracker.add('ExternalIngestClient::deleteFilesetByName'),
471478
token
472479
);
473480
const requestId = resp.headers.get('x-github-request-id');
@@ -476,7 +483,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
476483
this.logService.info(`ExternalIngestClient::deleteFilesetByName(): Deleted: ${fileSetName}`);
477484
}
478485

479-
async searchFilesets(filesetName: string, rootUri: URI, prompt: string, limit: number, token: CancellationToken): Promise<CodeSearchResult> {
486+
async searchFilesets(filesetName: string, rootUri: URI, prompt: string, limit: number, callTracker: CallTracker, token: CancellationToken): Promise<CodeSearchResult> {
480487
const authToken = await this.getAuthToken();
481488
if (!authToken) {
482489
this.logService.warn('ExternalIngestClient::searchFilesets(): No auth token available');
@@ -490,7 +497,7 @@ export class ExternalIngestClient extends Disposable implements IExternalIngestC
490497
scoping_query: `fileset:${filesetName}`,
491498
embedding_model: embeddingType.id,
492499
limit,
493-
}, {}, token);
500+
}, {}, callTracker.add('ExternalIngestClient::searchFilesets'), token);
494501

495502
const body = await resp.json() as SearchFilesetsResponse;
496503
return {

‎src/platform/workspaceChunkSearch/node/codeSearch/externalIngestIndex.ts‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import * as l10n from '@vscode/l10n';
88
import * as fs from 'node:fs';
99
import sql from 'node:sqlite';
1010
import { Result } from '../../../../util/common/result';
11+
import { CallTracker } from '../../../../util/common/telemetryCorrelationId';
1112
import { Limiter, raceCancellationError } from '../../../../util/vs/base/common/async';
1213
import { CancellationToken } from '../../../../util/vs/base/common/cancellation';
1314
import { Emitter } from '../../../../util/vs/base/common/event';
@@ -181,7 +182,7 @@ export class ExternalIngestIndex extends Disposable {
181182
* This deletes the remote file set and the checkpoint. We keep around the local database because it
182183
* has a cache of file shas.
183184
*/
184-
public async deleteIndex(token: CancellationToken): Promise<void> {
185+
public async deleteIndex(callTracker: CallTracker, token: CancellationToken): Promise<void> {
185186
const workspaceFolders = this._workspaceService.getWorkspaceFolders();
186187
if (!workspaceFolders.length) {
187188
return;
@@ -192,7 +193,7 @@ export class ExternalIngestIndex extends Disposable {
192193
this._logService.info(`ExternalIngestIndex: Deleting index for fileset ${filesetName}`);
193194

194195
try {
195-
await this._client.deleteFileset(filesetName, token);
196+
await this._client.deleteFileset(filesetName, callTracker, token);
196197
this.clearCurrentIndexCheckpoint();
197198
this._onDidChangeState.fire();
198199

@@ -251,7 +252,7 @@ export class ExternalIngestIndex extends Disposable {
251252
return this._initializePromise;
252253
}
253254

254-
async doIngest(onProgress: (message: string) => void, token: CancellationToken): Promise<Result<true, TriggerIndexingError>> {
255+
async doIngest(callTracker: CallTracker, onProgress: (message: string) => void, token: CancellationToken): Promise<Result<true, TriggerIndexingError>> {
255256
await raceCancellationError(this.initialize(), token);
256257

257258
const workspaceFolders = this._workspaceService.getWorkspaceFolders();
@@ -288,6 +289,7 @@ export class ExternalIngestIndex extends Disposable {
288289
this.getFilesetName(primaryRoot),
289290
currentCheckpoint,
290291
this.getFilesToIndexFromDb(token),
292+
callTracker,
291293
token,
292294
wrappedOnProgress
293295
);
@@ -350,7 +352,7 @@ export class ExternalIngestIndex extends Disposable {
350352
return updatePromise;
351353
}
352354

353-
async search(sizing: StrategySearchSizing, query: WorkspaceChunkQueryWithEmbeddings, token: CancellationToken): Promise<readonly FileChunkAndScore[]> {
355+
async search(sizing: StrategySearchSizing, query: WorkspaceChunkQueryWithEmbeddings, callTracker: CallTracker, token: CancellationToken): Promise<readonly FileChunkAndScore[]> {
354356
const workspaceFolders = this._workspaceService.getWorkspaceFolders();
355357
if (!workspaceFolders.length) {
356358
return [];
@@ -361,7 +363,7 @@ export class ExternalIngestIndex extends Disposable {
361363
try {
362364
const resolvedQuery = await query.resolveQuery(token);
363365

364-
await raceCancellationError(this.doIngest(() => { }, token), token);
366+
await raceCancellationError(this.doIngest(callTracker, () => { }, token), token);
365367

366368
// TODO: search changed files too
367369
const primaryRoot = workspaceFolders[0];
@@ -370,6 +372,7 @@ export class ExternalIngestIndex extends Disposable {
370372
primaryRoot,
371373
resolvedQuery,
372374
sizing.maxResultCountHint,
375+
callTracker,
373376
token), token);
374377

375378
/* __GDPR__

‎src/platform/workspaceChunkSearch/node/workspaceChunkSearchService.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import type * as vscode from 'vscode';
88
import { createFencedCodeBlock, getLanguageId } from '../../../util/common/markdown';
99
import { Result } from '../../../util/common/result';
1010
import { createServiceIdentifier } from '../../../util/common/services';
11-
import { TelemetryCorrelationId } from '../../../util/common/telemetryCorrelationId';
11+
import { CallTracker, TelemetryCorrelationId } from '../../../util/common/telemetryCorrelationId';
1212
import { TokenizerType } from '../../../util/common/tokenizer';
1313
import { coalesce } from '../../../util/vs/base/common/arrays';
1414
import { CancelablePromise, createCancelablePromise, raceCancellationError, raceTimeout } from '../../../util/vs/base/common/async';
@@ -349,7 +349,9 @@ class WorkspaceChunkSearchServiceImpl extends Disposable implements IWorkspaceCh
349349
}
350350

351351
deleteExternalIngestWorkspaceIndex(): Promise<void> {
352-
return this._codeSearchChunkSearch.deleteExternalIngestWorkspaceIndex(CancellationToken.None);
352+
return this._codeSearchChunkSearch.deleteExternalIngestWorkspaceIndex(
353+
new CallTracker('WorkspaceChunkSearchService::deleteExternalIngestWorkspaceIndex'),
354+
CancellationToken.None);
353355
}
354356

355357
async searchFileChunks(

0 commit comments

Comments
 (0)
⚡