Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/cli/src/ui/AppContainer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -694,6 +694,7 @@ export const AppContainer = (props: AppContainerProps) => {
settings.setValue(scope, 'security.auth.selectedType', authType);

try {
config.setRemoteAdminSettings(undefined);
await config.refreshAuth(authType);
setAuthState(AuthState.Authenticated);
} catch (e) {
Expand Down
94 changes: 54 additions & 40 deletions packages/core/src/code_assist/admin/admin_controls.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,7 @@ describe('Admin Controls', () => {
// Should still start polling
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: true,
adminControlsApplicable: true,
});
await vi.advanceTimersByTimeAsync(5 * 60 * 1000);

Expand All @@ -363,7 +364,10 @@ describe('Admin Controls', () => {
});

it('should fetch from server if no cachedSettings provided', async () => {
const serverResponse = { strictModeDisabled: false };
const serverResponse = {
strictModeDisabled: false,
adminControlsApplicable: true,
};
(mockServer.fetchAdminControls as Mock).mockResolvedValue(serverResponse);

const result = await fetchAdminControls(
Expand All @@ -386,31 +390,24 @@ describe('Admin Controls', () => {
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);
});

it('should return empty object on fetch error and still start polling', async () => {
(mockServer.fetchAdminControls as Mock).mockRejectedValue(
new Error('Network error'),
);
const result = await fetchAdminControls(
mockServer,
undefined,
true,
mockOnSettingsChanged,
);
it('should throw error on fetch error and NOT start polling', async () => {
const error = new Error('Network error');
(mockServer.fetchAdminControls as Mock).mockRejectedValue(error);

expect(result).toEqual({});
await expect(
fetchAdminControls(mockServer, undefined, true, mockOnSettingsChanged),
).rejects.toThrow(error);

// Polling should have been started and should retry
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: false,
});
// Polling should NOT have been started
// Advance timers just to be absolutely sure
await vi.advanceTimersByTimeAsync(5 * 60 * 1000);
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(2); // Initial + poll
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); // Only initial fetch
});

it('should return empty object on 403 fetch error and STOP polling', async () => {
const error403 = new Error('Forbidden');
Object.assign(error403, { status: 403 });
(mockServer.fetchAdminControls as Mock).mockRejectedValue(error403);
it('should return empty object on adminControlsApplicable false and STOP polling', async () => {
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
adminControlsApplicable: false,
});

const result = await fetchAdminControls(
mockServer,
Expand All @@ -421,7 +418,7 @@ describe('Admin Controls', () => {

expect(result).toEqual({});

// Advance time - should NOT poll because of 403
// Advance time - should NOT poll because of adminControlsApplicable: false
await vi.advanceTimersByTimeAsync(5 * 60 * 1000);
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1); // Only the initial call
});
Expand All @@ -430,6 +427,7 @@ describe('Admin Controls', () => {
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: false,
unknownField: 'bad',
adminControlsApplicable: true,
});

const result = await fetchAdminControls(
Expand All @@ -455,7 +453,9 @@ describe('Admin Controls', () => {
});

it('should reset polling interval if called again', async () => {
(mockServer.fetchAdminControls as Mock).mockResolvedValue({});
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
adminControlsApplicable: true,
});

// First call
await fetchAdminControls(
Expand Down Expand Up @@ -514,6 +514,7 @@ describe('Admin Controls', () => {
const serverResponse = {
strictModeDisabled: true,
unknownField: 'should be removed',
adminControlsApplicable: true,
};
(mockServer.fetchAdminControls as Mock).mockResolvedValue(serverResponse);

Expand All @@ -532,30 +533,32 @@ describe('Admin Controls', () => {
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);
});

it('should return empty object on 403 fetch error', async () => {
const error403 = new Error('Forbidden');
Object.assign(error403, { status: 403 });
(mockServer.fetchAdminControls as Mock).mockRejectedValue(error403);
it('should return empty object on adminControlsApplicable false', async () => {
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
adminControlsApplicable: false,
});

const result = await fetchAdminControlsOnce(mockServer, true);
expect(result).toEqual({});
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);
});

it('should return empty object on any other fetch error', async () => {
(mockServer.fetchAdminControls as Mock).mockRejectedValue(
new Error('Network error'),
it('should throw error on any other fetch error', async () => {
const error = new Error('Network error');
(mockServer.fetchAdminControls as Mock).mockRejectedValue(error);
await expect(fetchAdminControlsOnce(mockServer, true)).rejects.toThrow(
error,
);
const result = await fetchAdminControlsOnce(mockServer, true);
expect(result).toEqual({});
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);
});

it('should not start or stop any polling timers', async () => {
const setIntervalSpy = vi.spyOn(global, 'setInterval');
const clearIntervalSpy = vi.spyOn(global, 'clearInterval');

(mockServer.fetchAdminControls as Mock).mockResolvedValue({});
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
adminControlsApplicable: true,
});
await fetchAdminControlsOnce(mockServer, true);

expect(setIntervalSpy).not.toHaveBeenCalled();
Expand All @@ -568,6 +571,7 @@ describe('Admin Controls', () => {
// Initial fetch
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: true,
adminControlsApplicable: true,
});
await fetchAdminControls(
mockServer,
Expand All @@ -579,6 +583,7 @@ describe('Admin Controls', () => {
// Update for next poll
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: false,
adminControlsApplicable: true,
});

// Fast forward
Expand All @@ -598,7 +603,10 @@ describe('Admin Controls', () => {
});

it('should NOT emit if settings are deeply equal but not the same instance', async () => {
const settings = { strictModeDisabled: false };
const settings = {
strictModeDisabled: false,
adminControlsApplicable: true,
};
(mockServer.fetchAdminControls as Mock).mockResolvedValue(settings);

await fetchAdminControls(
Expand All @@ -613,6 +621,7 @@ describe('Admin Controls', () => {
// Next poll returns a different object with the same values
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: false,
adminControlsApplicable: true,
});
await vi.advanceTimersByTimeAsync(5 * 60 * 1000);

Expand All @@ -623,6 +632,7 @@ describe('Admin Controls', () => {
// Initial fetch is successful
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: true,
adminControlsApplicable: true,
});
await fetchAdminControls(
mockServer,
Expand All @@ -643,6 +653,7 @@ describe('Admin Controls', () => {
// Subsequent poll succeeds with new data
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: false,
adminControlsApplicable: true,
});
await vi.advanceTimersByTimeAsync(5 * 60 * 1000);
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(3);
Expand All @@ -659,10 +670,11 @@ describe('Admin Controls', () => {
});
});

it('should STOP polling if server returns 403', async () => {
it('should STOP polling if server returns adminControlsApplicable false', async () => {
// Initial fetch is successful
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
strictModeDisabled: true,
adminControlsApplicable: true,
});
await fetchAdminControls(
mockServer,
Expand All @@ -672,10 +684,10 @@ describe('Admin Controls', () => {
);
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(1);

// Next poll returns 403
const error403 = new Error('Forbidden');
Object.assign(error403, { status: 403 });
(mockServer.fetchAdminControls as Mock).mockRejectedValue(error403);
// Next poll returns adminControlsApplicable: false
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
adminControlsApplicable: false,
});

await vi.advanceTimersByTimeAsync(5 * 60 * 1000);
expect(mockServer.fetchAdminControls).toHaveBeenCalledTimes(2);
Expand All @@ -688,7 +700,9 @@ describe('Admin Controls', () => {

describe('stopAdminControlsPolling', () => {
it('should stop polling after it has started', async () => {
(mockServer.fetchAdminControls as Mock).mockResolvedValue({});
(mockServer.fetchAdminControls as Mock).mockResolvedValue({
adminControlsApplicable: true,
});

// Start polling
await fetchAdminControls(
Expand Down
51 changes: 21 additions & 30 deletions packages/core/src/code_assist/admin/admin_controls.ts
Original file line number Diff line number Diff line change
Expand Up @@ -80,15 +80,6 @@ export function sanitizeAdminSettings(
};
}

function isGaxiosError(error: unknown): error is { status: number } {
return (
typeof error === 'object' &&
error !== null &&
'status' in error &&
typeof (error as { status: unknown }).status === 'number'
);
}

/**
* Fetches the admin controls from the server if enabled by experiment flag.
* Safely handles polling start/stop based on the flag and server availability.
Expand Down Expand Up @@ -123,22 +114,20 @@ export async function fetchAdminControls(
const rawSettings = await server.fetchAdminControls({
project: server.projectId,
});

if (rawSettings.adminControlsApplicable !== true) {
stopAdminControlsPolling();
currentSettings = undefined;
return {};
}

const sanitizedSettings = sanitizeAdminSettings(rawSettings);
currentSettings = sanitizedSettings;
startAdminControlsPolling(server, server.projectId, onSettingsChanged);
return sanitizedSettings;
} catch (e) {
// Non-enterprise users don't have access to fetch settings.
if (isGaxiosError(e) && e.status === 403) {
stopAdminControlsPolling();
currentSettings = undefined;
return {};
}
debugLogger.error('Failed to fetch admin controls: ', e);
Comment thread
skeshive marked this conversation as resolved.
Resolved
// If initial fetch fails, start polling to retry.
currentSettings = {};
startAdminControlsPolling(server, server.projectId, onSettingsChanged);
return {};
throw e;
}
}

Expand All @@ -162,17 +151,18 @@ export async function fetchAdminControlsOnce(
const rawSettings = await server.fetchAdminControls({
project: server.projectId,
});
return sanitizeAdminSettings(rawSettings);
} catch (e) {
// Non-enterprise users don't have access to fetch settings.
if (isGaxiosError(e) && e.status === 403) {

if (rawSettings.adminControlsApplicable !== true) {
return {};
}

return sanitizeAdminSettings(rawSettings);
} catch (e) {
debugLogger.error(
'Failed to fetch admin controls: ',
e instanceof Error ? e.message : e,
);
return {};
throw e;
}
}

Expand All @@ -192,19 +182,20 @@ function startAdminControlsPolling(
const rawSettings = await server.fetchAdminControls({
project,
});

if (rawSettings.adminControlsApplicable !== true) {
stopAdminControlsPolling();
currentSettings = undefined;
return;
}

const newSettings = sanitizeAdminSettings(rawSettings);

if (!isDeepStrictEqual(newSettings, currentSettings)) {
currentSettings = newSettings;
onSettingsChanged(newSettings);
}
} catch (e) {
// Non-enterprise users don't have access to fetch settings.
if (isGaxiosError(e) && e.status === 403) {
stopAdminControlsPolling();
currentSettings = undefined;
return;
}
debugLogger.error('Failed to poll admin controls: ', e);
}
},
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/code_assist/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -355,4 +355,5 @@ export const FetchAdminControlsResponseSchema = z.object({
strictModeDisabled: z.boolean().optional(),
mcpSetting: McpSettingSchema.optional(),
cliFeatureSetting: CliFeatureSettingSchema.optional(),
adminControlsApplicable: z.boolean().optional(),
});
2 changes: 1 addition & 1 deletion packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1153,7 +1153,7 @@ export class Config {
return this.remoteAdminSettings;
}

setRemoteAdminSettings(settings: AdminControlsSettings): void {
setRemoteAdminSettings(settings: AdminControlsSettings | undefined): void {
this.remoteAdminSettings = settings;
}

Expand Down
Loading