Skip to content

Commit 42a0f56

Browse files
chrmartiCopilot
chrmarti
and
Copilot
committed
Add INESProvider.updateTreatmentVariables() (#1932)
* Add INESProvider.updateTreatmentVariables() * Update src/lib/node/chatLibMain.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Dispose --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 27aa2ca commit 42a0f56

File tree

2 files changed

+273
-2
lines changed

2 files changed

+273
-2
lines changed

‎chat-lib/test/simpleExperimentationService.spec.ts‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { describe, expect, it } from 'vitest';
7+
import { SimpleExperimentationService } from '../src/main';
8+
9+
describe('SimpleExperimentationService', () => {
10+
it('should initialize with no treatment variables', async () => {
11+
const service = new SimpleExperimentationService(false);
12+
await service.hasTreatments();
13+
14+
expect(service.getTreatmentVariable('nonexistent')).toBeUndefined();
15+
16+
service.dispose();
17+
});
18+
19+
it('should update multiple treatment variables at once', () => {
20+
const service = new SimpleExperimentationService(false);
21+
const variables: Record<string, boolean | number | string> = {
22+
'feature-a': true,
23+
'feature-b': false,
24+
'max-count': 100,
25+
'experiment-id': 'exp-123'
26+
};
27+
28+
service.updateTreatmentVariables(variables);
29+
30+
expect(service.getTreatmentVariable<boolean>('feature-a')).toBe(true);
31+
expect(service.getTreatmentVariable<boolean>('feature-b')).toBe(false);
32+
expect(service.getTreatmentVariable<number>('max-count')).toBe(100);
33+
expect(service.getTreatmentVariable<string>('experiment-id')).toBe('exp-123');
34+
35+
service.dispose();
36+
});
37+
38+
it('should fire onDidTreatmentsChange event with all changed variables', () => {
39+
const service = new SimpleExperimentationService(false);
40+
const events: string[][] = [];
41+
42+
service.onDidTreatmentsChange((event) => {
43+
events.push(event.affectedTreatmentVariables);
44+
});
45+
46+
const variables: Record<string, boolean | number | string> = {
47+
'feature-a': true,
48+
'feature-b': 50,
49+
'feature-c': 'test'
50+
};
51+
service.updateTreatmentVariables(variables);
52+
53+
expect(events).toHaveLength(1);
54+
expect(events[0]).toHaveLength(3);
55+
expect(events[0]).toContain('feature-a');
56+
expect(events[0]).toContain('feature-b');
57+
expect(events[0]).toContain('feature-c');
58+
59+
service.dispose();
60+
});
61+
62+
it('should not fire onDidTreatmentsChange event when no variables change', () => {
63+
const service = new SimpleExperimentationService(false);
64+
const events: string[][] = [];
65+
66+
// Set initial value
67+
const variables1: Record<string, boolean | number | string> = {
68+
'feature-flag': true
69+
};
70+
service.updateTreatmentVariables(variables1);
71+
72+
service.onDidTreatmentsChange((event) => {
73+
events.push(event.affectedTreatmentVariables);
74+
});
75+
76+
// Update with same value
77+
const variables2: Record<string, boolean | number | string> = {
78+
'feature-flag': true
79+
};
80+
service.updateTreatmentVariables(variables2);
81+
82+
expect(events).toHaveLength(0);
83+
84+
service.dispose();
85+
});
86+
87+
it('should fire onDidTreatmentsChange event only for changed variables', () => {
88+
const service = new SimpleExperimentationService(false);
89+
90+
// Set initial values
91+
const variables1: Record<string, boolean | number | string> = {
92+
'feature-a': true,
93+
'feature-b': 50
94+
};
95+
service.updateTreatmentVariables(variables1);
96+
97+
const events: string[][] = [];
98+
service.onDidTreatmentsChange((event) => {
99+
events.push(event.affectedTreatmentVariables);
100+
});
101+
102+
// Update with one changed value and one unchanged
103+
const variables2: Record<string, boolean | number | string> = {
104+
'feature-a': true, // unchanged
105+
'feature-b': 100 // changed
106+
};
107+
service.updateTreatmentVariables(variables2);
108+
109+
expect(events).toHaveLength(1);
110+
expect(events[0]).toEqual(['feature-b']);
111+
112+
service.dispose();
113+
});
114+
115+
it('should overwrite existing treatment variables', () => {
116+
const service = new SimpleExperimentationService(false);
117+
118+
const variables1: Record<string, boolean | number | string> = {
119+
'feature-flag': true
120+
};
121+
service.updateTreatmentVariables(variables1);
122+
123+
expect(service.getTreatmentVariable<boolean>('feature-flag')).toBe(true);
124+
125+
const variables2: Record<string, boolean | number | string> = {
126+
'feature-flag': false
127+
};
128+
service.updateTreatmentVariables(variables2);
129+
130+
expect(service.getTreatmentVariable<boolean>('feature-flag')).toBe(false);
131+
132+
service.dispose();
133+
});
134+
135+
it('should wait for treatment variables when waitForTreatmentVariables = true', async () => {
136+
const service = new SimpleExperimentationService(true);
137+
138+
// hasTreatments() should not resolve until updateTreatmentVariables() is called
139+
let hasTreatmentsResolved = false;
140+
const hasTreatmentsPromise = service.hasTreatments().then(() => {
141+
hasTreatmentsResolved = true;
142+
});
143+
144+
// Give a bit of time to ensure promise doesn't resolve immediately
145+
await new Promise(resolve => setTimeout(resolve, 10));
146+
expect(hasTreatmentsResolved).toBe(false);
147+
148+
// Now update treatment variables
149+
const variables: Record<string, boolean | number | string> = {
150+
'test-feature': true
151+
};
152+
service.updateTreatmentVariables(variables);
153+
154+
// Wait for hasTreatments to resolve
155+
await hasTreatmentsPromise;
156+
expect(hasTreatmentsResolved).toBe(true);
157+
expect(service.getTreatmentVariable<boolean>('test-feature')).toBe(true);
158+
159+
service.dispose();
160+
});
161+
162+
it('should remove treatment variable when omitted from update', () => {
163+
const service = new SimpleExperimentationService(false);
164+
165+
// Set initial variables
166+
const variables1: Record<string, boolean | number | string> = {
167+
'feature-a': true,
168+
'feature-b': 50,
169+
'feature-c': 'test'
170+
};
171+
service.updateTreatmentVariables(variables1);
172+
173+
expect(service.getTreatmentVariable<boolean>('feature-a')).toBe(true);
174+
expect(service.getTreatmentVariable<number>('feature-b')).toBe(50);
175+
expect(service.getTreatmentVariable<string>('feature-c')).toBe('test');
176+
177+
const events: string[][] = [];
178+
service.onDidTreatmentsChange((event) => {
179+
events.push(event.affectedTreatmentVariables);
180+
});
181+
182+
// Remove feature-b by omitting it
183+
const variables2: Record<string, boolean | number | string> = {
184+
'feature-a': true,
185+
'feature-c': 'test'
186+
};
187+
service.updateTreatmentVariables(variables2);
188+
189+
expect(service.getTreatmentVariable<boolean>('feature-a')).toBe(true);
190+
expect(service.getTreatmentVariable<number>('feature-b')).toBeUndefined();
191+
expect(service.getTreatmentVariable<string>('feature-c')).toBe('test');
192+
193+
expect(events).toHaveLength(1);
194+
expect(events[0]).toEqual(['feature-b']);
195+
196+
service.dispose();
197+
});
198+
});

‎src/lib/node/chatLibMain.ts‎

Lines changed: 75 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ import { NullRequestLogger } from '../../platform/requestLogger/node/nullRequest
5050
import { IRequestLogger } from '../../platform/requestLogger/node/requestLogger';
5151
import { ISimulationTestContext, NulSimulationTestContext } from '../../platform/simulationTestContext/common/simulationTestContext';
5252
import { ISnippyService, NullSnippyService } from '../../platform/snippy/common/snippyService';
53-
import { IExperimentationService, NullExperimentationService } from '../../platform/telemetry/common/nullExperimentationService';
53+
import { IExperimentationService, TreatmentsChangeEvent } from '../../platform/telemetry/common/nullExperimentationService';
5454
import { ITelemetryService, TelemetryDestination, TelemetryEventMeasurements, TelemetryEventProperties } from '../../platform/telemetry/common/telemetry';
5555
import { eventPropertiesToSimpleObject } from '../../platform/telemetry/common/telemetryData';
5656
import { ITokenizerProvider, TokenizerProvider } from '../../platform/tokenizer/node/tokenizer';
@@ -61,6 +61,7 @@ import { Disposable } from '../../util/vs/base/common/lifecycle';
6161
import { generateUuid } from '../../util/vs/base/common/uuid';
6262
import { SyncDescriptor } from '../../util/vs/platform/instantiation/common/descriptors';
6363
import { IInstantiationService } from '../../util/vs/platform/instantiation/common/instantiation';
64+
import { Emitter } from '../../util/vs/base/common/event';
6465

6566
/**
6667
* Log levels (taken from vscode.d.ts)
@@ -113,6 +114,11 @@ export interface INESProviderOptions {
113114
readonly copilotTokenManager: ICopilotTokenManager;
114115
readonly telemetrySender: ITelemetrySender;
115116
readonly logTarget?: ILogTarget;
117+
/**
118+
* If true, the provider will wait for treatment variables to be set.
119+
* INESProvider.updateTreatmentVariables() must be called to unblock.
120+
*/
121+
readonly waitForTreatmentVariables?: boolean;
116122
}
117123

118124
export interface INESResult {
@@ -132,6 +138,7 @@ export interface INESProvider<T extends INESResult = INESResult> {
132138
handleAcceptance(suggestion: T): void;
133139
handleRejection(suggestion: T): void;
134140
handleIgnored(suggestion: T, supersededByRequestUuid: T | undefined): void;
141+
updateTreatmentVariables(variables: Record<string, boolean | number | string>): void;
135142
dispose(): void;
136143
}
137144

@@ -264,13 +271,20 @@ class NESProvider extends Disposable implements INESProvider<NESResult> {
264271
throw e;
265272
}
266273
}
274+
275+
updateTreatmentVariables(variables: Record<string, boolean | number | string>) {
276+
if (this._expService instanceof SimpleExperimentationService) {
277+
this._expService.updateTreatmentVariables(variables);
278+
}
279+
}
280+
267281
}
268282

269283
function setupServices(options: INESProviderOptions) {
270284
const { fetcher, copilotTokenManager, telemetrySender, logTarget } = options;
271285
const builder = new InstantiationServiceBuilder();
272286
builder.define(IConfigurationService, new SyncDescriptor(DefaultsOnlyConfigurationService));
273-
builder.define(IExperimentationService, new SyncDescriptor(NullExperimentationService));
287+
builder.define(IExperimentationService, new SyncDescriptor(SimpleExperimentationService, [options.waitForTreatmentVariables]));
274288
builder.define(ISimulationTestContext, new SyncDescriptor(NulSimulationTestContext));
275289
builder.define(IWorkspaceService, new SyncDescriptor(NullWorkspaceService));
276290
builder.define(IDiffService, new SyncDescriptor(DiffServiceImpl, [false]));
@@ -303,6 +317,65 @@ function setupServices(options: INESProviderOptions) {
303317
return builder.seal();
304318
}
305319

320+
export class SimpleExperimentationService extends Disposable implements IExperimentationService {
321+
322+
declare readonly _serviceBrand: undefined;
323+
324+
private readonly variables: Record<string, boolean | number | string> = {};
325+
private readonly _onDidTreatmentsChange = this._register(new Emitter<TreatmentsChangeEvent>());
326+
readonly onDidTreatmentsChange = this._onDidTreatmentsChange.event;
327+
328+
private readonly waitFor: Promise<void>;
329+
private readonly resolveWaitFor: () => void;
330+
331+
constructor(
332+
waitForTreatmentVariables: boolean | undefined,
333+
) {
334+
super();
335+
if (waitForTreatmentVariables) {
336+
let resolveWaitFor: () => void;
337+
this.waitFor = new Promise<void>(resolve => {
338+
resolveWaitFor = resolve;
339+
});
340+
this.resolveWaitFor = resolveWaitFor!;
341+
} else {
342+
this.waitFor = Promise.resolve();
343+
this.resolveWaitFor = () => { };
344+
}
345+
}
346+
347+
async hasTreatments(): Promise<void> {
348+
return this.waitFor;
349+
}
350+
351+
getTreatmentVariable<T extends boolean | number | string>(name: string): T | undefined {
352+
return this.variables[name] as T | undefined;
353+
}
354+
355+
async setCompletionsFilters(_filters: Map<string, string>): Promise<void> { }
356+
357+
updateTreatmentVariables(variables: Record<string, boolean | number | string>): void {
358+
const changedVariables: string[] = [];
359+
for (const [key, value] of Object.entries(variables)) {
360+
const existing = this.variables[key];
361+
if (existing !== value) {
362+
this.variables[key] = value;
363+
changedVariables.push(key);
364+
}
365+
}
366+
for (const key of Object.keys(this.variables)) {
367+
if (!Object.hasOwn(variables, key)) {
368+
delete this.variables[key];
369+
changedVariables.push(key);
370+
}
371+
}
372+
if (changedVariables.length > 0) {
373+
this._onDidTreatmentsChange.fire({ affectedTreatmentVariables: changedVariables });
374+
}
375+
this.resolveWaitFor();
376+
}
377+
}
378+
306379
class SingleFetcherService implements IFetcherService {
307380

308381
declare readonly _serviceBrand: undefined;

0 commit comments

Comments
 (0)