Cleanup extension update logic (#10514)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Jacob MacDonald
2025-10-03 21:06:26 -07:00
committed by GitHub
parent 1a06282061
commit 7f8537a130
10 changed files with 310 additions and 292 deletions

View File

@@ -5,43 +5,23 @@
*/
import type { GeminiCLIExtension } from '@google/gemini-cli-core';
import {
updateAllUpdatableExtensions,
updateExtension,
} from '../../config/extensions/update.js';
import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
import { MessageType } from '../types.js';
import { extensionsCommand } from './extensionsCommand.js';
import { type CommandContext } from './types.js';
import {
describe,
it,
expect,
vi,
beforeEach,
type MockedFunction,
} from 'vitest';
import { ExtensionUpdateState } from '../state/extensions.js';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { type ExtensionUpdateAction } from '../state/extensions.js';
vi.mock('../../config/extensions/update.js', () => ({
updateExtension: vi.fn(),
updateAllUpdatableExtensions: vi.fn(),
checkForAllExtensionUpdates: vi.fn(),
}));
const mockUpdateExtension = updateExtension as MockedFunction<
typeof updateExtension
>;
const mockUpdateAllUpdatableExtensions =
updateAllUpdatableExtensions as MockedFunction<
typeof updateAllUpdatableExtensions
>;
const mockGetExtensions = vi.fn();
describe('extensionsCommand', () => {
let mockContext: CommandContext;
const mockDispatchExtensionState = vi.fn();
beforeEach(() => {
vi.resetAllMocks();
@@ -53,7 +33,7 @@ describe('extensionsCommand', () => {
},
},
ui: {
dispatchExtensionStateUpdate: vi.fn(),
dispatchExtensionStateUpdate: mockDispatchExtensionState,
},
});
});
@@ -93,7 +73,14 @@ describe('extensionsCommand', () => {
});
it('should inform user if there are no extensions to update with --all', async () => {
mockUpdateAllUpdatableExtensions.mockResolvedValue([]);
mockDispatchExtensionState.mockImplementationOnce(
(action: ExtensionUpdateAction) => {
if (action.type === 'SCHEDULE_UPDATE') {
action.payload.onComplete([]);
}
},
);
await updateAction(mockContext, '--all');
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
{
@@ -105,18 +92,24 @@ describe('extensionsCommand', () => {
});
it('should call setPendingItem and addItem in a finally block on success', async () => {
mockUpdateAllUpdatableExtensions.mockResolvedValue([
{
name: 'ext-one',
originalVersion: '1.0.0',
updatedVersion: '1.0.1',
mockDispatchExtensionState.mockImplementationOnce(
(action: ExtensionUpdateAction) => {
if (action.type === 'SCHEDULE_UPDATE') {
action.payload.onComplete([
{
name: 'ext-one',
originalVersion: '1.0.0',
updatedVersion: '1.0.1',
},
{
name: 'ext-two',
originalVersion: '2.0.0',
updatedVersion: '2.0.1',
},
]);
}
},
{
name: 'ext-two',
originalVersion: '2.0.0',
updatedVersion: '2.0.1',
},
]);
);
await updateAction(mockContext, '--all');
expect(mockContext.ui.setPendingItem).toHaveBeenCalledWith({
type: MessageType.EXTENSIONS_LIST,
@@ -131,9 +124,9 @@ describe('extensionsCommand', () => {
});
it('should call setPendingItem and addItem in a finally block on failure', async () => {
mockUpdateAllUpdatableExtensions.mockRejectedValue(
new Error('Something went wrong'),
);
mockDispatchExtensionState.mockImplementationOnce((_) => {
throw new Error('Something went wrong');
});
await updateAction(mockContext, '--all');
expect(mockContext.ui.setPendingItem).toHaveBeenCalledWith({
type: MessageType.EXTENSIONS_LIST,
@@ -155,95 +148,58 @@ describe('extensionsCommand', () => {
});
it('should update a single extension by name', async () => {
const extension: GeminiCLIExtension = {
name: 'ext-one',
version: '1.0.0',
isActive: true,
path: '/test/dir/ext-one',
installMetadata: {
type: 'git',
autoUpdate: false,
source: 'https://github.com/some/extension.git',
mockDispatchExtensionState.mockImplementationOnce(
(action: ExtensionUpdateAction) => {
if (action.type === 'SCHEDULE_UPDATE') {
action.payload.onComplete([
{
name: 'ext-one',
originalVersion: '1.0.0',
updatedVersion: '1.0.1',
},
]);
}
},
};
mockUpdateExtension.mockResolvedValue({
name: extension.name,
originalVersion: extension.version,
updatedVersion: '1.0.1',
});
mockGetExtensions.mockReturnValue([extension]);
mockContext.ui.extensionsUpdateState.set(extension.name, {
status: ExtensionUpdateState.UPDATE_AVAILABLE,
processed: false,
});
await updateAction(mockContext, 'ext-one');
expect(mockUpdateExtension).toHaveBeenCalledWith(
extension,
'/test/dir',
expect.any(Function),
ExtensionUpdateState.UPDATE_AVAILABLE,
expect.any(Function),
);
});
it('should handle errors when updating a single extension', async () => {
mockUpdateExtension.mockRejectedValue(new Error('Extension not found'));
mockGetExtensions.mockReturnValue([]);
await updateAction(mockContext, 'ext-one');
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
{
type: MessageType.ERROR,
text: 'Extension ext-one not found.',
expect(mockDispatchExtensionState).toHaveBeenCalledWith({
type: 'SCHEDULE_UPDATE',
payload: {
all: false,
names: ['ext-one'],
onComplete: expect.any(Function),
},
expect.any(Number),
);
});
});
it('should update multiple extensions by name', async () => {
const extensionOne: GeminiCLIExtension = {
name: 'ext-one',
version: '1.0.0',
isActive: true,
path: '/test/dir/ext-one',
installMetadata: {
type: 'git',
autoUpdate: false,
source: 'https://github.com/some/extension.git',
mockDispatchExtensionState.mockImplementationOnce(
(action: ExtensionUpdateAction) => {
if (action.type === 'SCHEDULE_UPDATE') {
action.payload.onComplete([
{
name: 'ext-one',
originalVersion: '1.0.0',
updatedVersion: '1.0.1',
},
{
name: 'ext-two',
originalVersion: '1.0.0',
updatedVersion: '1.0.1',
},
]);
}
},
};
const extensionTwo: GeminiCLIExtension = {
name: 'ext-two',
version: '1.0.0',
isActive: true,
path: '/test/dir/ext-two',
installMetadata: {
type: 'git',
autoUpdate: false,
source: 'https://github.com/some/extension.git',
},
};
mockGetExtensions.mockReturnValue([extensionOne, extensionTwo]);
mockContext.ui.extensionsUpdateState.set(
extensionOne.name,
ExtensionUpdateState.UPDATE_AVAILABLE,
);
mockContext.ui.extensionsUpdateState.set(
extensionTwo.name,
ExtensionUpdateState.UPDATE_AVAILABLE,
);
mockUpdateExtension
.mockResolvedValueOnce({
name: 'ext-one',
originalVersion: '1.0.0',
updatedVersion: '1.0.1',
})
.mockResolvedValueOnce({
name: 'ext-two',
originalVersion: '2.0.0',
updatedVersion: '2.0.1',
});
await updateAction(mockContext, 'ext-one ext-two');
expect(mockUpdateExtension).toHaveBeenCalledTimes(2);
expect(mockDispatchExtensionState).toHaveBeenCalledWith({
type: 'SCHEDULE_UPDATE',
payload: {
all: false,
names: ['ext-one', 'ext-two'],
onComplete: expect.any(Function),
},
});
expect(mockContext.ui.setPendingItem).toHaveBeenCalledWith({
type: MessageType.EXTENSIONS_LIST,
});

View File

@@ -4,15 +4,8 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { requestConsentInteractive } from '../../config/extension.js';
import {
updateAllUpdatableExtensions,
type ExtensionUpdateInfo,
updateExtension,
checkForAllExtensionUpdates,
} from '../../config/extensions/update.js';
import type { ExtensionUpdateInfo } from '../../config/extension.js';
import { getErrorMessage } from '../../utils/errors.js';
import { ExtensionUpdateState } from '../state/extensions.js';
import { MessageType } from '../types.js';
import {
type CommandContext,
@@ -29,11 +22,10 @@ async function listAction(context: CommandContext) {
);
}
async function updateAction(context: CommandContext, args: string) {
function updateAction(context: CommandContext, args: string): Promise<void> {
const updateArgs = args.split(' ').filter((value) => value.length > 0);
const all = updateArgs.length === 1 && updateArgs[0] === '--all';
const names = all ? undefined : updateArgs;
let updateInfos: ExtensionUpdateInfo[] = [];
const names = all ? null : updateArgs;
if (!all && names?.length === 0) {
context.ui.addItem(
@@ -43,32 +35,48 @@ async function updateAction(context: CommandContext, args: string) {
},
Date.now(),
);
return;
return Promise.resolve();
}
try {
await checkForAllExtensionUpdates(
context.services.config!.getExtensions(),
context.ui.dispatchExtensionStateUpdate,
let resolveUpdateComplete: (updateInfo: ExtensionUpdateInfo[]) => void;
const updateComplete = new Promise<ExtensionUpdateInfo[]>(
(resolve) => (resolveUpdateComplete = resolve),
);
updateComplete.then((updateInfos) => {
if (updateInfos.length === 0) {
context.ui.addItem(
{
type: MessageType.INFO,
text: 'No extensions to update.',
},
Date.now(),
);
}
context.ui.addItem(
{
type: MessageType.EXTENSIONS_LIST,
},
Date.now(),
);
context.ui.setPendingItem(null);
});
try {
context.ui.setPendingItem({
type: MessageType.EXTENSIONS_LIST,
});
if (all) {
updateInfos = await updateAllUpdatableExtensions(
context.services.config!.getWorkingDir(),
// We don't have the ability to prompt for consent yet in this flow.
(description) =>
requestConsentInteractive(
description,
context.ui.addConfirmUpdateExtensionRequest,
),
context.services.config!.getExtensions(),
context.ui.extensionsUpdateState,
context.ui.dispatchExtensionStateUpdate,
);
} else if (names?.length) {
const workingDir = context.services.config!.getWorkingDir();
context.ui.dispatchExtensionStateUpdate({
type: 'SCHEDULE_UPDATE',
payload: {
all,
names,
onComplete: (updateInfos) => {
resolveUpdateComplete(updateInfos);
},
},
});
if (names?.length) {
const extensions = context.services.config!.getExtensions();
for (const name of names) {
const extension = extensions.find(
@@ -84,33 +92,10 @@ async function updateAction(context: CommandContext, args: string) {
);
continue;
}
const updateInfo = await updateExtension(
extension,
workingDir,
(description) =>
requestConsentInteractive(
description,
context.ui.addConfirmUpdateExtensionRequest,
),
context.ui.extensionsUpdateState.get(extension.name)?.status ??
ExtensionUpdateState.UNKNOWN,
context.ui.dispatchExtensionStateUpdate,
);
if (updateInfo) updateInfos.push(updateInfo);
}
}
if (updateInfos.length === 0) {
context.ui.addItem(
{
type: MessageType.INFO,
text: 'No extensions to update.',
},
Date.now(),
);
return;
}
} catch (error) {
resolveUpdateComplete!([]);
context.ui.addItem(
{
type: MessageType.ERROR,
@@ -118,15 +103,8 @@ async function updateAction(context: CommandContext, args: string) {
},
Date.now(),
);
} finally {
context.ui.addItem(
{
type: MessageType.EXTENSIONS_LIST,
},
Date.now(),
);
context.ui.setPendingItem(null);
}
return updateComplete.then((_) => {});
}
const listExtensionsCommand: SlashCommand = {

View File

@@ -76,7 +76,7 @@ describe('useExtensionUpdates', () => {
const cwd = '/test/cwd';
vi.mocked(checkForAllExtensionUpdates).mockImplementation(
async (extensions, dispatch) => {
async (_extensions, dispatch, _cwd) => {
dispatch({
type: 'SET_STATE',
payload: {
@@ -122,7 +122,7 @@ describe('useExtensionUpdates', () => {
const addItem = vi.fn();
vi.mocked(checkForAllExtensionUpdates).mockImplementation(
async (extensions, dispatch) => {
async (_extensions, dispatch, _cwd) => {
dispatch({
type: 'SET_STATE',
payload: {
@@ -195,7 +195,7 @@ describe('useExtensionUpdates', () => {
const addItem = vi.fn();
vi.mocked(checkForAllExtensionUpdates).mockImplementation(
async (extensions, dispatch) => {
async (_extensions, dispatch, _cwd) => {
dispatch({
type: 'SET_STATE',
payload: {
@@ -280,7 +280,7 @@ describe('useExtensionUpdates', () => {
const cwd = '/test/cwd';
vi.mocked(checkForAllExtensionUpdates).mockImplementation(
async (extensions, dispatch) => {
async (_extensions, dispatch, _cwd) => {
dispatch({ type: 'BATCH_CHECK_START' });
dispatch({
type: 'SET_STATE',

View File

@@ -18,7 +18,10 @@ import {
checkForAllExtensionUpdates,
updateExtension,
} from '../../config/extensions/update.js';
import { requestConsentInteractive } from '../../config/extension.js';
import {
requestConsentInteractive,
type ExtensionUpdateInfo,
} from '../../config/extension.js';
import { checkExhaustive } from '../../utils/checks.js';
type ConfirmationRequestWrapper = {
@@ -41,7 +44,6 @@ function confirmationRequestsReducer(
return state.filter((r) => r !== action.request);
default:
checkExhaustive(action);
return state;
}
}
@@ -80,40 +82,77 @@ export const useExtensionUpdates = (
);
useEffect(() => {
(async () => {
await checkForAllExtensionUpdates(
extensions,
dispatchExtensionStateUpdate,
const extensionsToCheck = extensions.filter((extension) => {
const currentStatus = extensionsUpdateState.extensionStatuses.get(
extension.name,
);
})();
}, [extensions, extensions.length, dispatchExtensionStateUpdate]);
if (!currentStatus) return true;
const currentState = currentStatus.status;
return !currentState || currentState === ExtensionUpdateState.UNKNOWN;
});
if (extensionsToCheck.length === 0) return;
checkForAllExtensionUpdates(
extensionsToCheck,
dispatchExtensionStateUpdate,
cwd,
);
}, [
extensions,
extensionsUpdateState.extensionStatuses,
cwd,
dispatchExtensionStateUpdate,
]);
useEffect(() => {
if (extensionsUpdateState.batchChecksInProgress > 0) {
return;
}
const scheduledUpdate = extensionsUpdateState.scheduledUpdate;
if (scheduledUpdate) {
dispatchExtensionStateUpdate({
type: 'CLEAR_SCHEDULED_UPDATE',
});
}
function shouldDoUpdate(extension: GeminiCLIExtension): boolean {
if (scheduledUpdate) {
if (scheduledUpdate.all) {
return true;
}
return scheduledUpdate.names?.includes(extension.name) === true;
} else {
return extension.installMetadata?.autoUpdate === true;
}
}
let extensionsWithUpdatesCount = 0;
// We only notify if we have unprocessed extensions in the UPDATE_AVAILABLE
// state.
let shouldNotifyOfUpdates = false;
const updatePromises: Array<Promise<ExtensionUpdateInfo | undefined>> = [];
for (const extension of extensions) {
const currentState = extensionsUpdateState.extensionStatuses.get(
extension.name,
);
if (
!currentState ||
currentState.processed ||
currentState.status !== ExtensionUpdateState.UPDATE_AVAILABLE
) {
continue;
}
// Mark as processed immediately to avoid re-triggering.
dispatchExtensionStateUpdate({
type: 'SET_PROCESSED',
payload: { name: extension.name, processed: true },
});
if (extension.installMetadata?.autoUpdate) {
updateExtension(
const shouldUpdate = shouldDoUpdate(extension);
if (!shouldUpdate) {
extensionsWithUpdatesCount++;
if (!currentState.notified) {
// Mark as processed immediately to avoid re-triggering.
dispatchExtensionStateUpdate({
type: 'SET_NOTIFIED',
payload: { name: extension.name, notified: true },
});
shouldNotifyOfUpdates = true;
}
} else {
const updatePromise = updateExtension(
extension,
cwd,
(description) =>
@@ -123,7 +162,9 @@ export const useExtensionUpdates = (
),
currentState.status,
dispatchExtensionStateUpdate,
)
);
updatePromises.push(updatePromise);
updatePromise
.then((result) => {
if (!result) return;
addItem(
@@ -143,11 +184,9 @@ export const useExtensionUpdates = (
Date.now(),
);
});
} else {
extensionsWithUpdatesCount++;
}
}
if (extensionsWithUpdatesCount > 0) {
if (shouldNotifyOfUpdates) {
const s = extensionsWithUpdatesCount > 1 ? 's' : '';
addItem(
{
@@ -157,6 +196,18 @@ export const useExtensionUpdates = (
Date.now(),
);
}
if (scheduledUpdate) {
Promise.all(updatePromises).then((results) => {
const nonNullResults = results.filter((result) => result != null);
scheduledUpdate.onCompleteCallbacks.forEach((callback) => {
try {
callback(nonNullResults);
} catch (e) {
console.error(getErrorMessage(e));
}
});
});
}
}, [
extensions,
extensionsUpdateState,

View File

@@ -4,6 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
import type { ExtensionUpdateInfo } from '../../config/extension.js';
import { checkExhaustive } from '../../utils/checks.js';
export enum ExtensionUpdateState {
@@ -19,17 +20,34 @@ export enum ExtensionUpdateState {
export interface ExtensionUpdateStatus {
status: ExtensionUpdateState;
processed: boolean;
notified: boolean;
}
export interface ExtensionUpdatesState {
extensionStatuses: Map<string, ExtensionUpdateStatus>;
batchChecksInProgress: number;
// Explicitly scheduled updates.
scheduledUpdate: ScheduledUpdate | null;
}
export interface ScheduledUpdate {
names: string[] | null;
all: boolean;
onCompleteCallbacks: OnCompleteUpdate[];
}
export interface ScheduleUpdateArgs {
names: string[] | null;
all: boolean;
onComplete: OnCompleteUpdate;
}
type OnCompleteUpdate = (updateInfos: ExtensionUpdateInfo[]) => void;
export const initialExtensionUpdatesState: ExtensionUpdatesState = {
extensionStatuses: new Map(),
batchChecksInProgress: 0,
scheduledUpdate: null,
};
export type ExtensionUpdateAction =
@@ -38,11 +56,13 @@ export type ExtensionUpdateAction =
payload: { name: string; state: ExtensionUpdateState };
}
| {
type: 'SET_PROCESSED';
payload: { name: string; processed: boolean };
type: 'SET_NOTIFIED';
payload: { name: string; notified: boolean };
}
| { type: 'BATCH_CHECK_START' }
| { type: 'BATCH_CHECK_END' };
| { type: 'BATCH_CHECK_END' }
| { type: 'SCHEDULE_UPDATE'; payload: ScheduleUpdateArgs }
| { type: 'CLEAR_SCHEDULED_UPDATE' };
export function extensionUpdatesReducer(
state: ExtensionUpdatesState,
@@ -57,19 +77,19 @@ export function extensionUpdatesReducer(
const newStatuses = new Map(state.extensionStatuses);
newStatuses.set(action.payload.name, {
status: action.payload.state,
processed: false,
notified: false,
});
return { ...state, extensionStatuses: newStatuses };
}
case 'SET_PROCESSED': {
case 'SET_NOTIFIED': {
const existing = state.extensionStatuses.get(action.payload.name);
if (!existing || existing.processed === action.payload.processed) {
if (!existing || existing.notified === action.payload.notified) {
return state;
}
const newStatuses = new Map(state.extensionStatuses);
newStatuses.set(action.payload.name, {
...existing,
processed: action.payload.processed,
notified: action.payload.notified,
});
return { ...state, extensionStatuses: newStatuses };
}
@@ -83,6 +103,27 @@ export function extensionUpdatesReducer(
...state,
batchChecksInProgress: state.batchChecksInProgress - 1,
};
case 'SCHEDULE_UPDATE':
return {
...state,
// If there is a pre-existing scheduled update, we merge them.
scheduledUpdate: {
all: state.scheduledUpdate?.all || action.payload.all,
names: [
...(state.scheduledUpdate?.names ?? []),
...(action.payload.names ?? []),
],
onCompleteCallbacks: [
...(state.scheduledUpdate?.onCompleteCallbacks ?? []),
action.payload.onComplete,
],
},
};
case 'CLEAR_SCHEDULED_UPDATE':
return {
...state,
scheduledUpdate: null,
};
default:
checkExhaustive(action);
}