Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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
8 changes: 5 additions & 3 deletions packages/a2a-server/src/agent/task.ts
Original file line number Diff line number Diff line change
Expand Up @@ -131,9 +131,11 @@ export class Task {
this.autoExecute = autoExecute;
this.config.setFallbackModelHandler(
// For a2a-server, we want to automatically switch to the fallback model
// and retry the current request seamlessly. The 'retry_always' intent
// achieves this, ensuring a smooth fallback experience for the user.
async () => 'retry_always',
// for future requests without retrying the current one.
async (failedModel, fallbackModel) => {
this.config.activateFallbackMode(fallbackModel, failedModel);
return 'stop';
},
);
}

Expand Down
3 changes: 3 additions & 0 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1938,6 +1938,9 @@
}

activateFallbackMode(model: string, failedModel?: string): void {
debugLogger.log(
`Model fallback activated: switching from ${failedModel ?? 'unknown'} to ${model}`,
);
if (this.getActiveModel() !== model) {
this.setModel(model, true);
}
Expand Down Expand Up @@ -3571,9 +3574,9 @@
// Gemini API key users should have the ability to manually select the
// old preview flash model.
if (authType === AuthType.USE_GEMINI) {
setFlashModels('gemini-3-flash-preview', 'gemini-3.5-flash');

Check warning on line 3577 in packages/core/src/config/config.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.5". Please make sure this change is appropriate to submit.
} else {
setFlashModels('gemini-3.5-flash', 'gemini-3.5-flash');

Check warning on line 3579 in packages/core/src/config/config.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.5". Please make sure this change is appropriate to submit.

Check warning on line 3579 in packages/core/src/config/config.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.5". Please make sure this change is appropriate to submit.
}
} else {
setFlashModels('gemini-3-flash-preview', 'gemini-2.5-flash');
Expand Down
30 changes: 30 additions & 0 deletions packages/core/src/utils/flashFallback.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,36 @@ describe('Retry Utility Fallback Integration', () => {
expect(mockApiCall).toHaveBeenCalledTimes(3);
});

it('should call onPersistent429 immediately on attempt 1 when classifyGoogleError returns TerminalQuotaError', async () => {
const mockApiCall = vi
.fn()
.mockRejectedValue(
new TerminalQuotaError('Capacity exhausted', mockGoogleApiError),
);

const mockPersistent429Callback = vi.fn(
async () =>
// Return null to stop retrying after fallback attempt
null,
);

const promise = retryWithBackoff(mockApiCall, {
maxAttempts: 10, // High maxAttempts to prove we don't wait for max attempts
initialDelayMs: 1,
maxDelayMs: 10,
onPersistent429: mockPersistent429Callback,
authType: AuthType.LOGIN_WITH_GOOGLE,
});

await expect(promise).rejects.toThrow('Capacity exhausted');
expect(mockApiCall).toHaveBeenCalledTimes(1); // Only called once because it's terminal and fallback returned null
expect(mockPersistent429Callback).toHaveBeenCalledTimes(1);
expect(mockPersistent429Callback).toHaveBeenCalledWith(
AuthType.LOGIN_WITH_GOOGLE,
expect.any(TerminalQuotaError),
);
});

it('should trigger onPersistent429 when HTTP 499 persists through all retry attempts', async () => {
let fallbackCalled = false;
const mockError: HttpError = new Error('Simulated 499 error');
Expand Down
45 changes: 42 additions & 3 deletions packages/core/src/utils/googleQuotaErrors.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ describe('classifyGoogleError', () => {
}
});

it('should return RetryableQuotaError with delay for 503 Service Unavailable with RetryInfo', () => {
it('should return TerminalQuotaError for MODEL_CAPACITY_EXHAUSTED even with RetryInfo headers', () => {
const apiError: GoogleApiError = {
code: 503,
message:
Expand All @@ -103,8 +103,7 @@ describe('classifyGoogleError', () => {
};
vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
const result = classifyGoogleError(new Error());
expect(result).toBeInstanceOf(RetryableQuotaError);
expect((result as RetryableQuotaError).retryDelayMs).toBe(9000);
expect(result).toBeInstanceOf(TerminalQuotaError);
});

it('should return TerminalQuotaError for MODEL_CAPACITY_EXHAUSTED when no retry delay is specified', () => {
Expand All @@ -126,6 +125,24 @@ describe('classifyGoogleError', () => {
expect(result).toBeInstanceOf(TerminalQuotaError);
});

it('should return TerminalQuotaError for structured error with details when message contains capacity exhaustion keywords', () => {
const apiError: GoogleApiError = {
code: 429,
message: 'You have exhausted your capacity on this model.',
details: [
{
'@type': 'type.googleapis.com/google.rpc.Help',
links: [
{ description: 'Learn more', url: 'https://support.google.com' },
],
},
],
};
vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
const result = classifyGoogleError(new Error());
expect(result).toBeInstanceOf(TerminalQuotaError);
});

it('should return TerminalQuotaError for MODEL_CAPACITY_EXHAUSTED even when the domain is not a Cloud Code domain (domain-agnostic)', () => {
const apiError: GoogleApiError = {
code: 429,
Expand Down Expand Up @@ -396,6 +413,28 @@ describe('classifyGoogleError', () => {
expect((result as TerminalQuotaError).reason).toBe('RATE_LIMIT_EXCEEDED');
});

it('should return TerminalQuotaError for Cloud Code RATE_LIMIT_EXCEEDED without a specified server delay', () => {
const apiError: GoogleApiError = {
code: 429,
message: 'Rate limit exceeded',
details: [
{
'@type': 'type.googleapis.com/google.rpc.ErrorInfo',
reason: 'RATE_LIMIT_EXCEEDED',
domain: 'cloudcode-pa.googleapis.com',
metadata: {
uiMessage: 'true',
model: 'gemini-2.5-pro',
},
},
],
};
vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
const result = classifyGoogleError(new Error());
expect(result).toBeInstanceOf(TerminalQuotaError);
expect((result as TerminalQuotaError).reason).toBe('RATE_LIMIT_EXCEEDED');
});

it('should return TerminalQuotaError for Cloud Code QUOTA_EXHAUSTED', () => {
const apiError: GoogleApiError = {
code: 429,
Expand Down
76 changes: 48 additions & 28 deletions packages/core/src/utils/googleQuotaErrors.ts
Original file line number Diff line number Diff line change
Expand Up @@ -289,16 +289,23 @@ export function classifyGoogleError(error: unknown): unknown {
return new RetryableQuotaError(errorMessage, cause, retryDelaySeconds);
}
} else if (status === 429 || status === 499 || status === 503) {
// Fallback: If it is a 429, 499, or 503 but doesn't have a specific "retry in" message,
// assume it is a temporary rate limit and retry.
return new RetryableQuotaError(
errorMessage,
googleApiError ?? {
code: status,
message: errorMessage,
details: [],
},
);
const cause = googleApiError ?? {
code: status,
message: errorMessage,
details: [],
};

// If the error message indicates capacity exhaustion, classify as TerminalQuotaError
if (
/exhausted your capacity|capacity exceeded|MODEL_CAPACITY_EXHAUSTED/i.test(
errorMessage,
)
) {
return new TerminalQuotaError(errorMessage, cause);
}

// Fallback: assume it is a temporary rate limit and retry.
return new RetryableQuotaError(errorMessage, cause);
}

return error; // Not a retryable error we can handle with structured details or a parsable retry message.
Expand Down Expand Up @@ -338,8 +345,11 @@ export function classifyGoogleError(error: unknown): unknown {
}

if (errorInfo) {
// INSUFFICIENT_G1_CREDITS_BALANCE is always terminal, regardless of domain
if (errorInfo.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE') {
// Always treat capacity exhaustion as terminal error to trigger immediate model fallback
if (
errorInfo.reason === 'MODEL_CAPACITY_EXHAUSTED' ||
errorInfo.reason === 'MODEL_CAPACITY_EXCEEDED'
) {
return new TerminalQuotaError(
googleApiError.message,
googleApiError,
Expand All @@ -348,28 +358,29 @@ export function classifyGoogleError(error: unknown): unknown {
);
}

if (
errorInfo.reason === 'MODEL_CAPACITY_EXHAUSTED' ||
errorInfo.reason === 'MODEL_CAPACITY_EXCEEDED'
) {
// If no server backoff delay is specified, treat capacity exhaustion as a terminal error
// to trigger immediate model fallback without retrying on the same exhausted model.
if (delaySeconds === undefined) {
return new TerminalQuotaError(
googleApiError.message,
googleApiError,
delaySeconds,
errorInfo.reason,
);
}
// Otherwise, fall through to RetryableQuotaError to honor the server's requested delay.
// INSUFFICIENT_G1_CREDITS_BALANCE is always terminal, regardless of domain
if (errorInfo.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE') {
return new TerminalQuotaError(
googleApiError.message,
googleApiError,
delaySeconds,
errorInfo.reason,
);
}

// New Cloud Code API quota handling
if (errorInfo.domain) {
if (isCloudCodeDomain(errorInfo.domain)) {
if (errorInfo.reason === 'RATE_LIMIT_EXCEEDED') {
const effectiveDelay = delaySeconds ?? 10;
if (delaySeconds === undefined) {
return new TerminalQuotaError(
googleApiError.message,
googleApiError,
undefined,
errorInfo.reason,
);
}
const effectiveDelay = delaySeconds;
if (effectiveDelay > MAX_RETRYABLE_DELAY_SECONDS) {
return new TerminalQuotaError(
googleApiError.message,
Expand Down Expand Up @@ -437,6 +448,15 @@ export function classifyGoogleError(error: unknown): unknown {
}
}

// If the error message indicates capacity exhaustion, classify as TerminalQuotaError
if (
/exhausted your capacity|capacity exceeded|MODEL_CAPACITY_EXHAUSTED/i.test(
errorMessage,
)
) {
return new TerminalQuotaError(errorMessage, googleApiError);
Comment thread
DavidAPierce marked this conversation as resolved.
}

// If we reached this point, the status is 429, 499, or 503 and we have details,
// but no specific violation was matched. We return a generic retryable error.
return new RetryableQuotaError(errorMessage, googleApiError);
Expand Down
Loading