feat(core): add agent decision mechanism for unretryable errors

This commit is contained in:
Sehoon Shon
2026-03-10 19:40:19 -04:00
parent e22d9917b7
commit ed7b9f7069
6 changed files with 270 additions and 79 deletions

View File

@@ -20,12 +20,14 @@ import { reportError } from '../utils/errorReporting.js';
import { getErrorMessage } from '../utils/errors.js';
import { logMalformedJsonResponse } from '../telemetry/loggers.js';
import { MalformedJsonResponseEvent, LlmRole } from '../telemetry/types.js';
import { retryWithBackoff } from '../utils/retry.js';
import { retryWithBackoff, type AgentDecision } from '../utils/retry.js';
import type { ModelConfigKey } from '../services/modelConfigService.js';
import {
applyModelSelection,
createAvailabilityContextProvider,
} from '../availability/policyHelpers.js';
import { PREVIEW_GEMINI_FLASH_MODEL } from '../config/models.js';
import { debugLogger } from '../utils/debugLogger.js';
const DEFAULT_MAX_ATTEMPTS = 5;
@@ -316,6 +318,92 @@ export class BaseLlmClient {
);
};
const onAgentDecisionCallback = async (
error: unknown,
attempt: number,
): Promise<AgentDecision> => {
try {
const errorMsg =
error instanceof Error ? error.message : String(error);
if (errorMsg.includes('429')) {
return 'stop';
}
const lastUserMessage = [...contents]
.reverse()
.find((c) => c.role === 'user');
const lastUserText = lastUserMessage?.parts
? lastUserMessage.parts.map((p) => p.text || '').join(' ')
: 'N/A';
const decisionPrompt = `You are a meta-agent deciding whether to retry an AI request that failed.
Error: ${errorMsg}
Attempt: ${attempt}
Target Model: ${currentModel}
Last User Request: ${lastUserText}
Based on the error, should we try to retry the exact same request?
Some errors are transient (e.g. network blips, internal server errors), others are terminal (e.g. safety blocks, invalid arguments).
Respond with a JSON object: {"action": "retry" | "stop"}`;
const decisionResponse = await this.contentGenerator.generateContent(
{
model: PREVIEW_GEMINI_FLASH_MODEL,
contents: [{ role: 'user', parts: [{ text: decisionPrompt }] }],
config: {
responseMimeType: 'application/json',
},
},
'agent-decision',
LlmRole.UTILITY_ROUTER,
);
const decisionText =
decisionResponse.candidates?.[0]?.content?.parts?.[0]?.text;
if (decisionText) {
const parsed = JSON.parse(decisionText) as unknown;
if (
typeof parsed === 'object' &&
parsed !== null &&
'action' in parsed
) {
const action = (parsed as { action: unknown }).action;
if (
action === 'retry' ||
action === 'stop' ||
action === 'modify_request'
) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return action as AgentDecision;
}
}
}
} catch (agentError) {
debugLogger.warn('Agent decision failed:', agentError);
}
return 'stop';
};
const onAgentDecisionWrapper = async (
error: unknown,
attempt: number,
): Promise<AgentDecision> => {
if (
typeof error === 'object' &&
error !== null &&
'message' in error &&
typeof (error as { message?: unknown }).message === 'string' &&
(error as { message?: unknown }).message === 'Agent Decision'
) {
return 'stop';
}
return onAgentDecisionCallback(error, attempt);
};
return await retryWithBackoff(apiCall, {
shouldRetryOnContent,
maxAttempts:
@@ -325,6 +413,7 @@ export class BaseLlmClient {
? (authType, error) =>
handleFallback(this.config, currentModel, authType, error)
: undefined,
onAgentDecision: onAgentDecisionWrapper,
authType:
this.authType ?? this.config.getContentGeneratorConfig()?.authType,
});

View File

@@ -19,13 +19,18 @@ import {
type GenerateContentParameters,
} from '@google/genai';
import { toParts } from '../code_assist/converter.js';
import { retryWithBackoff, isRetryableError } from '../utils/retry.js';
import {
retryWithBackoff,
isRetryableError,
type AgentDecision,
} from '../utils/retry.js';
import type { ValidationRequiredError } from '../utils/googleQuotaErrors.js';
import type { Config } from '../config/config.js';
import {
resolveModel,
isGemini2Model,
supportsModernFeatures,
PREVIEW_GEMINI_FLASH_MODEL,
} from '../config/models.js';
import { hasCycleInSchema } from '../tools/tools.js';
import type { StructuredError } from './turn.js';
@@ -41,7 +46,7 @@ import {
import {
ContentRetryEvent,
ContentRetryFailureEvent,
type LlmRole,
LlmRole,
} from '../telemetry/types.js';
import { handleFallback } from '../fallback/handler.js';
import { isFunctionResponse } from '../utils/messageInspectors.js';
@@ -53,6 +58,7 @@ import {
createAvailabilityContextProvider,
} from '../availability/policyHelpers.js';
import { coreEvents } from '../utils/events.js';
import { debugLogger } from '../utils/debugLogger.js';
export enum StreamEventType {
/** A regular content chunk from the API. */
@@ -624,9 +630,97 @@ export class GeminiChat {
);
};
const onAgentDecisionCallback = async (
error: unknown,
attempt: number,
): Promise<AgentDecision> => {
try {
const errorMsg = error instanceof Error ? error.message : String(error);
// Don't try to use the agent for 429s as it will likely also fail.
if (errorMsg.includes('429')) {
return 'stop';
}
const lastUserMessage = [...lastContentsToUse]
.reverse()
.find((c) => c.role === 'user');
const lastUserText = lastUserMessage?.parts
? lastUserMessage.parts.map((p) => p.text || '').join(' ')
: 'N/A';
const decisionPrompt = `You are a meta-agent deciding whether to retry an AI request that failed.
Error: ${errorMsg}
Attempt: ${attempt}
Target Model: ${lastModelToUse}
Last User Request: ${lastUserText}
Based on the error, should we try to retry the exact same request?
Some errors are transient (e.g. network blips, internal server errors), others are terminal (e.g. safety blocks, invalid arguments).
Respond with a JSON object: {"action": "retry" | "stop"}`;
const decisionResponse = await this.config
.getContentGenerator()
.generateContent(
{
model: PREVIEW_GEMINI_FLASH_MODEL,
contents: [{ role: 'user', parts: [{ text: decisionPrompt }] }],
config: {
responseMimeType: 'application/json',
},
},
'agent-decision',
LlmRole.UTILITY_ROUTER,
);
const decisionText =
decisionResponse.candidates?.[0]?.content?.parts?.[0]?.text;
if (decisionText) {
const parsed = JSON.parse(decisionText) as unknown;
if (
typeof parsed === 'object' &&
parsed !== null &&
'action' in parsed
) {
const action = (parsed as { action: unknown }).action;
if (
action === 'retry' ||
action === 'stop' ||
action === 'modify_request'
) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return action as AgentDecision;
}
}
}
} catch (agentError) {
debugLogger.warn('Agent decision failed:', agentError);
}
return 'stop';
};
const onAgentDecisionWrapper = async (
error: unknown,
attempt: number,
): Promise<AgentDecision> => {
if (
typeof error === 'object' &&
error !== null &&
'message' in error &&
typeof (error as { message?: unknown }).message === 'string' &&
(error as { message?: unknown }).message === 'Agent Decision'
) {
return 'stop';
}
return onAgentDecisionCallback(error, attempt);
};
const streamResponse = await retryWithBackoff(apiCall, {
onPersistent429: onPersistent429Callback,
onValidationRequired: onValidationRequiredCallback,
onAgentDecision: onAgentDecisionWrapper,
authType: this.config.getContentGeneratorConfig()?.authType,
retryFetchErrors: this.config.getRetryFetchErrors(),
signal: abortSignal,

View File

@@ -5,7 +5,6 @@
*/
import type { Config } from '../config/config.js';
import { AuthType } from '../core/contentGenerator.js';
import {
openBrowserSecurely,
shouldLaunchBrowser,
@@ -29,10 +28,6 @@ export async function handleFallback(
authType?: string,
error?: unknown,
): Promise<string | boolean | null> {
if (authType !== AuthType.LOGIN_WITH_GOOGLE) {
return null;
}
const chain = resolvePolicyChain(config);
const { failedPolicy, candidates } = buildFallbackPolicyContext(
chain,

View File

@@ -82,6 +82,35 @@ describe('retryWithBackoff', () => {
expect(mockFn).toHaveBeenCalledTimes(3);
});
it('should ask the agent when max attempts are reached, and retry if requested', async () => {
let calls = 0;
const mockFn = vi.fn(async () => {
calls++;
if (calls <= 2) {
throw new Error('Generic unretryable error');
}
return 'success';
});
const mockAgentDecision = vi.fn().mockResolvedValue('retry');
const promise = retryWithBackoff(mockFn, {
maxAttempts: 1, // Doesn't matter because the error is unretryable and will fail immediately, OR if retryable it will fail after max attempts
initialDelayMs: 10,
onAgentDecision: mockAgentDecision,
});
await vi.runAllTimersAsync();
const result = await promise;
expect(result).toBe('success');
expect(mockFn).toHaveBeenCalledTimes(3); // 1 original + 2 retries
expect(mockAgentDecision).toHaveBeenCalledTimes(2);
expect(mockAgentDecision).toHaveBeenCalledWith(
expect.any(Error),
1, // the attempt number when it failed
);
});
it('should throw an error if all attempts fail', async () => {
const mockFn = createFailingFunction(3);

View File

@@ -17,6 +17,12 @@ import { getErrorStatus, ModelNotFoundError } from './httpErrors.js';
import type { RetryAvailabilityContext } from '../availability/modelPolicy.js';
export type { RetryAvailabilityContext };
export type AgentDecision =
| 'retry'
| 'stop'
| { action: 'retry' | 'stop' | 'modify_request'; data?: unknown };
export const DEFAULT_MAX_ATTEMPTS = 10;
export interface RetryOptions {
@@ -32,6 +38,7 @@ export interface RetryOptions {
onValidationRequired?: (
error: ValidationRequiredError,
) => Promise<'verify' | 'change_auth' | 'cancel'>;
onAgentDecision?: (error: unknown, attempt: number) => Promise<AgentDecision>;
authType?: string;
retryFetchErrors?: boolean;
signal?: AbortSignal;
@@ -174,6 +181,7 @@ export async function retryWithBackoff<T>(
maxDelayMs,
onPersistent429,
onValidationRequired,
onAgentDecision,
authType,
shouldRetryOnError,
shouldRetryOnContent,
@@ -190,6 +198,29 @@ export async function retryWithBackoff<T>(
let attempt = 0;
let currentDelay = initialDelayMs;
const handleTerminalError = async (errToThrow: unknown) => {
if (onAgentDecision) {
try {
const decision = await onAgentDecision(errToThrow, attempt);
const decisionType =
typeof decision === 'string' ? decision : decision.action;
debugLogger.log(
`Agent evaluated error at attempt ${attempt}: ${decisionType}`,
);
if (decisionType === 'retry' || decisionType === 'modify_request') {
return true;
}
} catch (agentError) {
if (agentError !== errToThrow) {
debugLogger.warn('Agent decision failed:', agentError);
}
}
}
throw errToThrow;
};
while (attempt < maxAttempts) {
if (signal?.aborted) {
throw createAbortError();
@@ -248,7 +279,11 @@ export async function retryWithBackoff<T>(
}
}
// Terminal/not_found already recorded; nothing else to mark here.
throw classifiedError; // Throw if no fallback or fallback failed.
if (await handleTerminalError(classifiedError)) {
attempt = 0;
currentDelay = initialDelayMs;
continue;
}
}
// Handle ValidationRequiredError - user needs to verify before proceeding
@@ -268,7 +303,11 @@ export async function retryWithBackoff<T>(
debugLogger.warn('Validation handler failed:', validationError);
}
}
throw classifiedError;
if (await handleTerminalError(classifiedError)) {
attempt = 0;
currentDelay = initialDelayMs;
continue;
}
}
const is500 =
@@ -296,9 +335,15 @@ export async function retryWithBackoff<T>(
debugLogger.warn('Model fallback failed:', fallbackError);
}
}
throw classifiedError instanceof RetryableQuotaError
? classifiedError
: error;
const errToThrow =
classifiedError instanceof RetryableQuotaError
? classifiedError
: error;
if (await handleTerminalError(errToThrow)) {
attempt = 0;
currentDelay = initialDelayMs;
continue;
}
}
if (
@@ -340,7 +385,11 @@ export async function retryWithBackoff<T>(
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
!shouldRetryOnError(error as Error, retryFetchErrors)
) {
throw error;
if (await handleTerminalError(error)) {
attempt = 0;
currentDelay = initialDelayMs;
continue;
}
}
const errorStatus = getErrorStatus(error);

View File

@@ -1,65 +0,0 @@
## Summary
This PR implements a seamless migration path for extensions to move to a new
repository and optionally change their name without stranding existing users.
When an extension author sets the `migratedTo` field in their
`gemini-extension.json` and publishes an update to their old repository, the CLI
will detect this during the next update check. The CLI will then automatically
download the extension from the new repository, explicitly warn the user about
the migration (and any renaming) during the consent step, and seamlessly migrate
the installation and enablement status while cleaning up the old installation.
## Details
- **Configuration:** Added `migratedTo` property to `ExtensionConfig` and
`GeminiCLIExtension` to track the new repository URL.
- **Update checking & downloading:** Updated `checkForExtensionUpdate` and
`updateExtension` to inspect the `migratedTo` field. If present, the CLI
queries the new repository URL for an update and swaps the installation source
so the update resolves from the new location.
- **Migration & renaming logic (`ExtensionManager`):**
- `installOrUpdateExtension` now fully supports renaming. It transfers global
and workspace enablement states from the old extension name to the new one
and deletes the old extension directory.
- Added safeguards to block renaming if the new name conflicts with a
different, already-installed extension or if the destination directory
already exists.
- Exposed `getEnablementManager()` to `ExtensionManager` for better typing
during testing.
- **Consent messaging:** Refactored `maybeRequestConsentOrFail` to compute an
`isMigrating` flag (by detecting a change in the installation source). The
`extensionConsentString` output now explicitly informs users with messages
like: _"Migrating extension 'old-name' to a new repository, renaming to
'new-name', and installing updates."_
- **Documentation:** Documented the `migratedTo` field in
`docs/extensions/reference.md` and added a comprehensive guide in
`docs/extensions/releasing.md` explaining how extension maintainers can
transition users using this feature.
- **Testing:** Added extensive unit tests across `extension-manager.test.ts`,
`consent.test.ts`, `github.test.ts`, and `update.test.ts` to cover the new
migration and renaming logic.
## Related issues
N/A
## How to validate
1. **Unit tests:** Run all related tests to confirm everything passes:
```bash
npm run test -w @google/gemini-cli -- src/config/extensions/github.test.ts
npm run test -w @google/gemini-cli -- src/config/extensions/update.test.ts
npm run test -w @google/gemini-cli -- src/config/extensions/consent.test.ts
npm run test -w @google/gemini-cli -- src/config/extension-manager.test.ts
```
2. **End-to-end migration test:**
- Install a local or git extension.
- Update its `gemini-extension.json` to include a `migratedTo` field pointing
to a _different_ test repository.
- Run `gemini extensions check` to confirm it detects the update from the new
source.
- Run `gemini extensions update <extension>`.
- Verify that the consent prompt explicitly mentions the migration.
- Verify that the new extension is installed, the old directory is deleted,
and its enablement status carried over.