Fix internal-org and staff targeting in the new TAS endpoint (#339929)

Fix VS Code team targeting in the new TAS endpoint

extensions/copilot/src/platform/telemetry/test/node/assignmentFilters.spec.ts

@@ -93,6 +93,41 @@ describe('CopilotAssignmentsFilterProvider', () => {
 		]);
 	});
 
+	it('preserves VS Code team targeting and internal organization precedence across token changes', () => {
+		const { tokenStore, provider } = createProvider();
+		const organizations = [
+			['a5db0bcaae94032fe715fb34a5e4bce2'],
+			['4535c7beffc844b46bb1ed4aa04d759a'],
+			['551cca60ce19654d894e786220822482'],
+			['a5db0bcaae94032fe715fb34a5e4bce2', '4535c7beffc844b46bb1ed4aa04d759a', '551cca60ce19654d894e786220822482'],
+			['unrecognized-org'],
+			[],
+		];
+		const tokens = [
+			...organizations.map(organization_list => new CopilotToken(createTestExtendedTokenInfo({ organization_list }))),
+			new CopilotToken(createTestExtendedTokenInfo({ isVscodeTeamMember: true })),
+			undefined,
+		];
+
+		expect(tokens.map(token => {
+			tokenStore.copilotToken = token;
+			const filters = provider.getFilters();
+			return [
+				filters.get('github_core_ghmsftorexternal'),
+				filters.get('github_core_isghormsftstaff'),
+			];
+		})).toEqual([
+			['microsoft', '1'],
+			['github', '1'],
+			['vscode', '1'],
+			['vscode', '1'],
+			['external', '0'],
+			['external', '0'],
+			['vscode', '1'],
+			['external', '0'],
+		]);
+	});
+
 	it('preserves legacy CLS user-kind values and precedence', () => {
 		const { tokenStore, provider } = createProvider();
 		const knownOrganizations = [

extensions/copilot/src/platform/telemetry/vscode-node/microsoftExperimentationService.ts

@@ -254,7 +254,7 @@ export class CopilotAssignmentsFilterProvider implements IExperimentationFilterP
 		filters.set('github_core_organizationid', token?.organizationList.join(','));
 		filters.set('github_core_businessid', token?.enterpriseList.join(','));
 		filters.set('github_core_isghormsftstaff', internalOrg ? '1' : '0');
-		filters.set('github_core_ghmsftorexternal', internalOrg === 'github' ? 'github' : (internalOrg === 'microsoft' || internalOrg === 'vscode') ? 'microsoft' : 'external');
+		filters.set('github_core_ghmsftorexternal', internalOrg ?? 'external');
 		filters.set('github_core_userkind', token?.userKind || undefined);
 		filters.set('github_core_copilotsku', token?.sku);
 		filters.set('github_core_issn', token?.isSn() ? '1' : '0');

src/vs/workbench/services/assignment/common/assignmentFilters.ts

@@ -261,6 +261,7 @@ export class CopilotAssignmentFilterProvider extends Disposable implements IExpe
  */
 export enum GitHubAssignmentsFilter {
 	CopilotTrackingId = 'copilottrackingid',
+	/** Whether the account is staff or belongs to a recognized internal organization. */
 	IsGhOrMsftStaff = 'github_core_isghormsftstaff',
 	GhMsftOrExternal = 'github_core_ghmsftorexternal',
 }
@@ -272,6 +273,7 @@ export enum GitHubAssignmentsFilter {
 export class GitHubCoreAssignmentsFilterProvider extends Disposable implements IExperimentationFilterProvider {
 	private copilotTrackingId: string | undefined;
 	private internalOrg: 'vscode' | 'github' | 'microsoft' | undefined;
+	private isInternal = false;
 
 	private readonly _onDidChangeFilters = this._register(new Emitter<void>());
 	readonly onDidChangeFilters = this._onDidChangeFilters.event;
@@ -288,13 +290,15 @@ export class GitHubCoreAssignmentsFilterProvider extends Disposable implements I
 	private update(): void {
 		const newTrackingId = this._chatEntitlementService.copilotTrackingId ?? this.copilotTrackingId;
 		const newInternalOrg = getInternalOrg(this._chatEntitlementService.organisations);
+		const newIsInternal = this._chatEntitlementService.isInternal;
 
-		if (this.copilotTrackingId === newTrackingId && this.internalOrg === newInternalOrg) {
+		if (this.copilotTrackingId === newTrackingId && this.internalOrg === newInternalOrg && this.isInternal === newIsInternal) {
 			return;
 		}
 
 		this.copilotTrackingId = newTrackingId;
 		this.internalOrg = newInternalOrg;
+		this.isInternal = newIsInternal;
 
 		this._onDidChangeFilters.fire();
 	}
@@ -313,13 +317,9 @@ export class GitHubCoreAssignmentsFilterProvider extends Disposable implements I
 			case GitHubAssignmentsFilter.CopilotTrackingId:
 				return copilotTrackingId ?? null;
 			case GitHubAssignmentsFilter.IsGhOrMsftStaff:
-				return internalOrg ? '1' : '0';
+				return this._chatEntitlementService.isInternal ? '1' : '0';
 			case GitHubAssignmentsFilter.GhMsftOrExternal:
-				return internalOrg === 'github'
-					? 'github'
-					: (internalOrg === 'microsoft' || internalOrg === 'vscode')
-						? 'microsoft'
-						: 'external';
+				return internalOrg ?? 'external';
 			default:
 				return null;
 		}

src/vs/workbench/services/assignment/test/common/assignmentFilters.test.ts

@@ -8,6 +8,7 @@ import { ICopilotTokenInfo } from '../../../../../base/common/defaultAccount.js'
 import { Emitter, Event } from '../../../../../base/common/event.js';
 import { Lazy } from '../../../../../base/common/lazy.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js';
+import { isInternalAccount } from '../../../../../platform/assignment/common/assignment.js';
 import { IDefaultAccountService } from '../../../../../platform/defaultAccount/common/defaultAccount.js';
 import { IExtensionDescription } from '../../../../../platform/extensions/common/extensions.js';
 import { TestInstantiationService } from '../../../../../platform/instantiation/test/common/instantiationServiceMock.js';
@@ -154,6 +155,11 @@ suite('GitHubCoreAssignmentsFilterProvider', () => {
 
 		override copilotTrackingId: string | undefined;
 		override organisations: string[] | undefined;
+		isStaff: boolean | undefined;
+
+		override get isInternal(): boolean {
+			return isInternalAccount(this.isStaff, this.organisations);
+		}
 
 		fireChange(): void {
 			this._onDidChangeEntitlement.fire();
@@ -201,19 +207,62 @@ suite('GitHubCoreAssignmentsFilterProvider', () => {
 
 	test('reflects organization classification changes', () => {
 		const entitlement = new TestChatEntitlementService();
-		entitlement.organisations = undefined;
-
 		const provider = createProvider(entitlement);
+		const organizations = [
+			undefined,
+			['microsoft'],
+			['github'],
+			['Visual-Studio-Code'],
+			['microsoft', 'github', 'Visual-Studio-Code'],
+			['unrecognized-org'],
+			[],
+			undefined,
+		];
 
-		assert.strictEqual(provider.getFilterValue(GitHubAssignmentsFilter.IsGhOrMsftStaff), '0');
-		assert.strictEqual(provider.getFilterValue(GitHubAssignmentsFilter.GhMsftOrExternal), 'external');
+		assert.deepStrictEqual(organizations.map(organisations => {
+			entitlement.organisations = organisations;
+			entitlement.fireChange();
+			const filters = provider.getFilters();
+			return [
+				filters.get(GitHubAssignmentsFilter.GhMsftOrExternal),
+				filters.get(GitHubAssignmentsFilter.IsGhOrMsftStaff),
+			];
+		}), [
+			['external', '0'],
+			['microsoft', '1'],
+			['github', '1'],
+			['vscode', '1'],
+			['vscode', '1'],
+			['external', '0'],
+			['external', '0'],
+			['external', '0'],
+		]);
+	});
 
-		entitlement.organisations = ['microsoft'];
-		assert.strictEqual(provider.getFilterValue(GitHubAssignmentsFilter.IsGhOrMsftStaff), '1');
-		assert.strictEqual(provider.getFilterValue(GitHubAssignmentsFilter.GhMsftOrExternal), 'microsoft');
+	test('includes staff without internal organizations and notifies on staff-only changes', () => {
+		const entitlement = new TestChatEntitlementService();
+		entitlement.organisations = ['contoso'];
+		entitlement.isStaff = true;
+		const provider = createProvider(entitlement);
+		let changes = 0;
+		disposables.add(provider.onDidChangeFilters(() => changes++));
 
-		entitlement.organisations = ['Visual-Studio-Code'];
-		assert.strictEqual(provider.getFilterValue(GitHubAssignmentsFilter.GhMsftOrExternal), 'microsoft');
+		assert.deepStrictEqual([true, false, false, true, undefined].map(isStaff => {
+			entitlement.isStaff = isStaff;
+			entitlement.fireChange();
+			const filters = provider.getFilters();
+			return [
+				filters.get(GitHubAssignmentsFilter.IsGhOrMsftStaff),
+				filters.get(GitHubAssignmentsFilter.GhMsftOrExternal),
+				changes,
+			];
+		}), [
+			['1', 'external', 0],
+			['0', 'external', 1],
+			['0', 'external', 1],
+			['1', 'external', 2],
+			['0', 'external', 3],
+		]);
 	});
 
 	test('fires onDidChangeFilters only when relevant inputs change', () => {