Skip to content

Commit 39a7387

Browse files
lramos15lramos15
authored
Fix manage models being improperly shown (microsoft#261318)
1 parent 0b42c89 commit 39a7387

File tree

4 files changed

+53
-16
lines changed

4 files changed

+53
-16
lines changed

‎src/vs/workbench/contrib/chat/browser/actions/manageModelsActions.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { ThemeIcon } from '../../../../../base/common/themables.js';
1010
import { localize2 } from '../../../../../nls.js';
1111
import { Action2 } from '../../../../../platform/actions/common/actions.js';
1212
import { ICommandService } from '../../../../../platform/commands/common/commands.js';
13+
import { ContextKeyExpr } from '../../../../../platform/contextkey/common/contextkey.js';
1314
import { ServicesAccessor } from '../../../../../platform/instantiation/common/instantiation.js';
1415
import { IQuickInputService, IQuickPickItem } from '../../../../../platform/quickinput/common/quickInput.js';
1516
import { ChatContextKeys } from '../../common/chatContextKeys.js';
@@ -34,7 +35,12 @@ export class ManageModelsAction extends Action2 {
3435
id: ManageModelsAction.ID,
3536
title: localize2('manageLanguageModels', 'Manage Language Models...'),
3637
category: CHAT_CATEGORY,
37-
precondition: ChatContextKeys.enabled,
38+
precondition: ContextKeyExpr.and(ChatContextKeys.enabled, ContextKeyExpr.or(
39+
ChatContextKeys.Entitlement.free,
40+
ChatContextKeys.Entitlement.pro,
41+
ChatContextKeys.Entitlement.proPlus,
42+
ChatContextKeys.Entitlement.internal
43+
)),
3844
f1: true
3945
});
4046
}

‎src/vs/workbench/contrib/chat/browser/modelPicker/modelPickerActionItem.ts‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -53,17 +53,24 @@ function modelDelegateToWidgetActionsProvider(delegate: IModelPickerDelegate): I
5353
function getModelPickerActionBarActions(menuService: IMenuService, contextKeyService: IContextKeyService, commandService: ICommandService, chatEntitlementService: IChatEntitlementService): IAction[] {
5454
const additionalActions: IAction[] = [];
5555

56-
additionalActions.push({
57-
id: 'manageModels',
58-
label: localize('chat.manageModels', "Manage Models..."),
59-
enabled: true,
60-
tooltip: localize('chat.manageModels.tooltip', "Manage language models"),
61-
class: undefined,
62-
run: () => {
63-
const commandId = ManageModelsAction.ID;
64-
commandService.executeCommand(commandId);
65-
}
66-
});
56+
if (
57+
chatEntitlementService.entitlement === ChatEntitlement.Free ||
58+
chatEntitlementService.entitlement === ChatEntitlement.Pro ||
59+
chatEntitlementService.entitlement === ChatEntitlement.ProPlus ||
60+
chatEntitlementService.isInternal
61+
) {
62+
additionalActions.push({
63+
id: 'manageModels',
64+
label: localize('chat.manageModels', "Manage Models..."),
65+
enabled: true,
66+
tooltip: localize('chat.manageModels.tooltip', "Manage language models"),
67+
class: undefined,
68+
run: () => {
69+
const commandId = ManageModelsAction.ID;
70+
commandService.executeCommand(commandId);
71+
}
72+
});
73+
}
6774

6875
// Add upgrade option if entitlement is free
6976
if (chatEntitlementService.entitlement === ChatEntitlement.Free) {

‎src/vs/workbench/contrib/chat/common/chatContextKeys.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ export namespace ChatContextKeys {
7474
};
7575

7676
export const Entitlement = {
77+
internal: new RawContextKey<boolean>('chatEntitlementInternal', false, true), // True when user is a chat internal user.
7778
signedOut: new RawContextKey<boolean>('chatEntitlementSignedOut', false, true), // True when user is signed out.
7879
canSignUp: new RawContextKey<boolean>('chatPlanCanSignUp', false, true), // True when user can sign up to be a chat free user.
7980
free: new RawContextKey<boolean>('chatPlanFree', false, true), // True when user is a chat free user.

‎src/vs/workbench/contrib/chat/common/chatEntitlementService.ts‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ export interface IChatEntitlementService {
101101
readonly onDidChangeEntitlement: Event<void>;
102102

103103
readonly entitlement: ChatEntitlement;
104+
readonly isInternal: boolean;
104105

105106
readonly onDidChangeQuotaExceeded: Event<void>;
106107
readonly onDidChangeQuotaRemaining: Event<void>;
@@ -239,6 +240,10 @@ export class ChatEntitlementService extends Disposable implements IChatEntitleme
239240
return ChatEntitlement.Unresolved;
240241
}
241242

243+
get isInternal(): boolean {
244+
return this.contextKeyService.getContextKeyValue<boolean>(ChatContextKeys.Entitlement.internal.key) === true;
245+
}
246+
242247
//#endregion
243248

244249
//#region --- Quotas
@@ -384,6 +389,7 @@ interface IEntitlementsResponse extends ILegacyQuotaSnapshotResponse {
384389
readonly can_signup_for_limited: boolean;
385390
readonly chat_enabled: boolean;
386391
readonly copilot_plan: string;
392+
readonly organization_login_list: string[];
387393
readonly analytics_tracking_id: string;
388394
readonly limited_user_reset_date?: string; // for Copilot Free
389395
readonly quota_reset_date?: string; // for all other Copilot SKUs
@@ -396,6 +402,7 @@ interface IEntitlementsResponse extends ILegacyQuotaSnapshotResponse {
396402

397403
interface IEntitlements {
398404
readonly entitlement: ChatEntitlement;
405+
readonly isInternal?: boolean;
399406
readonly quotas?: IQuotas;
400407
}
401408

@@ -644,6 +651,7 @@ export class ChatEntitlementRequests extends Disposable {
644651

645652
const entitlements: IEntitlements = {
646653
entitlement,
654+
isInternal: this.containsInternalOrgs(entitlementsResponse.organization_login_list),
647655
quotas: this.toQuotas(entitlementsResponse)
648656
};
649657

@@ -718,6 +726,11 @@ export class ChatEntitlementRequests extends Disposable {
718726
return quotas;
719727
}
720728

729+
private containsInternalOrgs(organizationLogins: string[]): boolean {
730+
const internalOrgs = ['github', 'microsoft'];
731+
return organizationLogins.some(org => internalOrgs.includes(org));
732+
}
733+
721734
private async request(url: string, type: 'GET', body: undefined, sessions: AuthenticationSession[], token: CancellationToken): Promise<IRequestContext | undefined>;
722735
private async request(url: string, type: 'POST', body: object, sessions: AuthenticationSession[], token: CancellationToken): Promise<IRequestContext | undefined>;
723736
private async request(url: string, type: 'GET' | 'POST', body: object | undefined, sessions: AuthenticationSession[], token: CancellationToken): Promise<IRequestContext | undefined> {
@@ -759,7 +772,7 @@ export class ChatEntitlementRequests extends Disposable {
759772
private update(state: IEntitlements): void {
760773
this.state = state;
761774

762-
this.context.update({ entitlement: this.state.entitlement });
775+
this.context.update({ entitlement: this.state.entitlement, isInternal: this.state.isInternal });
763776

764777
if (state.quotas) {
765778
this.chatQuotasAccessor.acceptQuotas(state.quotas);
@@ -906,6 +919,11 @@ export interface IChatEntitlementContextState extends IChatSentiment {
906919
*/
907920
entitlement: ChatEntitlement;
908921

922+
/**
923+
* User is an internal chat user.
924+
*/
925+
isInternal: boolean;
926+
909927
/**
910928
* User is or was a registered Chat user.
911929
*/
@@ -925,6 +943,8 @@ export class ChatEntitlementContext extends Disposable {
925943
private readonly businessContextKey: IContextKey<boolean>;
926944
private readonly enterpriseContextKey: IContextKey<boolean>;
927945

946+
private readonly internalContextKey: IContextKey<boolean>;
947+
928948
private readonly hiddenContext: IContextKey<boolean>;
929949
private readonly laterContext: IContextKey<boolean>;
930950
private readonly installedContext: IContextKey<boolean>;
@@ -956,13 +976,14 @@ export class ChatEntitlementContext extends Disposable {
956976
this.proPlusContextKey = ChatContextKeys.Entitlement.proPlus.bindTo(contextKeyService);
957977
this.businessContextKey = ChatContextKeys.Entitlement.business.bindTo(contextKeyService);
958978
this.enterpriseContextKey = ChatContextKeys.Entitlement.enterprise.bindTo(contextKeyService);
979+
this.internalContextKey = ChatContextKeys.Entitlement.internal.bindTo(contextKeyService);
959980
this.hiddenContext = ChatContextKeys.Setup.hidden.bindTo(contextKeyService);
960981
this.laterContext = ChatContextKeys.Setup.later.bindTo(contextKeyService);
961982
this.installedContext = ChatContextKeys.Setup.installed.bindTo(contextKeyService);
962983
this.disabledContext = ChatContextKeys.Setup.disabled.bindTo(contextKeyService);
963984
this.untrustedContext = ChatContextKeys.Setup.untrusted.bindTo(contextKeyService);
964985

965-
this._state = this.storageService.getObject<IChatEntitlementContextState>(ChatEntitlementContext.CHAT_ENTITLEMENT_CONTEXT_STORAGE_KEY, StorageScope.PROFILE) ?? { entitlement: ChatEntitlement.Unknown };
986+
this._state = this.storageService.getObject<IChatEntitlementContextState>(ChatEntitlementContext.CHAT_ENTITLEMENT_CONTEXT_STORAGE_KEY, StorageScope.PROFILE) ?? { entitlement: ChatEntitlement.Unknown, isInternal: false };
966987

967988
this.checkExtensionInstallation();
968989
this.updateContextSync();
@@ -1004,9 +1025,10 @@ export class ChatEntitlementContext extends Disposable {
10041025
update(context: { installed: boolean; disabled: boolean; untrusted: boolean }): Promise<void>;
10051026
update(context: { hidden: boolean }): Promise<void>;
10061027
update(context: { later: boolean }): Promise<void>;
1007-
update(context: { entitlement: ChatEntitlement }): Promise<void>;
1008-
update(context: { installed?: boolean; disabled?: boolean; untrusted?: boolean; hidden?: boolean; later?: boolean; entitlement?: ChatEntitlement }): Promise<void> {
1028+
update(context: { entitlement: ChatEntitlement; isInternal?: boolean }): Promise<void>;
1029+
update(context: { installed?: boolean; disabled?: boolean; untrusted?: boolean; hidden?: boolean; later?: boolean; entitlement?: ChatEntitlement; isInternal?: boolean }): Promise<void> {
10091030
this.logService.trace(`[chat entitlement context] update(): ${JSON.stringify(context)}`);
1031+
this.state.isInternal = !!context.isInternal;
10101032

10111033
if (typeof context.installed === 'boolean' && typeof context.disabled === 'boolean' && typeof context.untrusted === 'boolean') {
10121034
this._state.installed = context.installed;
@@ -1060,6 +1082,7 @@ export class ChatEntitlementContext extends Disposable {
10601082
this.proPlusContextKey.set(this._state.entitlement === ChatEntitlement.ProPlus);
10611083
this.businessContextKey.set(this._state.entitlement === ChatEntitlement.Business);
10621084
this.enterpriseContextKey.set(this._state.entitlement === ChatEntitlement.Enterprise);
1085+
this.internalContextKey.set(this._state.isInternal);
10631086
this.hiddenContext.set(!!this._state.hidden);
10641087
this.laterContext.set(!!this._state.later);
10651088
this.installedContext.set(!!this._state.installed);

0 commit comments

Comments
 (0)