Skip to content

Commit 8671014

Browse files
committed
fix: add fetch retries for bearer path and extract token helper in e2e
Address PR review comments: - Add fetchWithRetry() for transient failure retries (429/5xx/network) on the MCP bearer-token HTTP path, matching the SDK retry behavior the SigV4 path gets for free - Extract duplicated Cognito token fetch into fetchCognitoAccessToken() helper in byo-custom-jwt e2e tests
1 parent 9ac6776 commit 8671014

2 files changed

Lines changed: 44 additions & 48 deletions

File tree

e2e-tests/byo-custom-jwt.test.ts

Lines changed: 20 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,23 @@ describe.sequential('e2e: BYO agent with CUSTOM_JWT auth', () => {
6767
const cognitoClient = new CognitoIdentityProviderClient({ region });
6868
const cfnClient = new CloudFormationClient({ region });
6969

70+
/** Fetch a Cognito access token via client_credentials flow. */
71+
async function fetchCognitoAccessToken(): Promise<string> {
72+
const tokenUrl = `https://${domainPrefix}.auth.${region}.amazoncognito.com/oauth2/token`;
73+
const tokenRes = await fetch(tokenUrl, {
74+
method: 'POST',
75+
headers: {
76+
'Content-Type': 'application/x-www-form-urlencoded',
77+
Authorization: `Basic ${Buffer.from(`${clientId}:${clientSecret}`).toString('base64')}`,
78+
},
79+
body: 'grant_type=client_credentials&scope=agentcore/invoke',
80+
});
81+
expect(tokenRes.ok, `Token fetch failed: ${tokenRes.status}`).toBe(true);
82+
const tokenJson = (await tokenRes.json()) as { access_token: string };
83+
expect(tokenJson.access_token, 'Should have received an access token').toBeTruthy();
84+
return tokenJson.access_token;
85+
}
86+
7087
beforeAll(async () => {
7188
if (!canRun) return;
7289

@@ -267,20 +284,7 @@ describe.sequential('e2e: BYO agent with CUSTOM_JWT auth', () => {
267284
async () => {
268285
expect(projectPath, 'Project should have been deployed').toBeTruthy();
269286

270-
// Fetch a Cognito access token via client_credentials flow
271-
const tokenUrl = `https://${domainPrefix}.auth.${region}.amazoncognito.com/oauth2/token`;
272-
const tokenRes = await fetch(tokenUrl, {
273-
method: 'POST',
274-
headers: {
275-
'Content-Type': 'application/x-www-form-urlencoded',
276-
Authorization: `Basic ${Buffer.from(`${clientId}:${clientSecret}`).toString('base64')}`,
277-
},
278-
body: 'grant_type=client_credentials&scope=agentcore/invoke',
279-
});
280-
expect(tokenRes.ok, `Token fetch failed: ${tokenRes.status}`).toBe(true);
281-
const tokenJson = (await tokenRes.json()) as { access_token: string };
282-
const accessToken = tokenJson.access_token;
283-
expect(accessToken, 'Should have received an access token').toBeTruthy();
287+
const accessToken = await fetchCognitoAccessToken();
284288

285289
// Invoke with bearer token — should NOT get auth mismatch
286290
const result = await runLocalCLI(
@@ -313,18 +317,7 @@ describe.sequential('e2e: BYO agent with CUSTOM_JWT auth', () => {
313317
async () => {
314318
expect(projectPath, 'Project should have been deployed').toBeTruthy();
315319

316-
const tokenUrl = `https://${domainPrefix}.auth.${region}.amazoncognito.com/oauth2/token`;
317-
const tokenRes = await fetch(tokenUrl, {
318-
method: 'POST',
319-
headers: {
320-
'Content-Type': 'application/x-www-form-urlencoded',
321-
Authorization: `Basic ${Buffer.from(`${clientId}:${clientSecret}`).toString('base64')}`,
322-
},
323-
body: 'grant_type=client_credentials&scope=agentcore/invoke',
324-
});
325-
expect(tokenRes.ok, `Token fetch failed: ${tokenRes.status}`).toBe(true);
326-
const tokenJson = (await tokenRes.json()) as { access_token: string };
327-
const accessToken = tokenJson.access_token;
320+
const accessToken = await fetchCognitoAccessToken();
328321

329322
const result = await runLocalCLI(
330323
['invoke', '--agent', mcpAgentName, 'list-tools', '--bearer-token', accessToken, '--json'],
@@ -342,18 +335,7 @@ describe.sequential('e2e: BYO agent with CUSTOM_JWT auth', () => {
342335
async () => {
343336
expect(projectPath, 'Project should have been deployed').toBeTruthy();
344337

345-
const tokenUrl = `https://${domainPrefix}.auth.${region}.amazoncognito.com/oauth2/token`;
346-
const tokenRes = await fetch(tokenUrl, {
347-
method: 'POST',
348-
headers: {
349-
'Content-Type': 'application/x-www-form-urlencoded',
350-
Authorization: `Basic ${Buffer.from(`${clientId}:${clientSecret}`).toString('base64')}`,
351-
},
352-
body: 'grant_type=client_credentials&scope=agentcore/invoke',
353-
});
354-
expect(tokenRes.ok, `Token fetch failed: ${tokenRes.status}`).toBe(true);
355-
const tokenJson = (await tokenRes.json()) as { access_token: string };
356-
const accessToken = tokenJson.access_token;
338+
const accessToken = await fetchCognitoAccessToken();
357339

358340
const result = await runLocalCLI(
359341
[

src/cli/aws/agentcore.ts

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -554,6 +554,28 @@ interface McpRpcResult {
554554
error?: { message?: string; code?: number };
555555
}
556556

557+
const TRANSIENT_STATUS_CODES = new Set([429, 500, 502, 503, 504]);
558+
const MAX_FETCH_RETRIES = 3;
559+
const FETCH_RETRY_DELAY_MS = 1000;
560+
561+
/** Retry-aware fetch for transient failures (5xx, 429, network errors). */
562+
async function fetchWithRetry(url: string, init: RequestInit, logger?: SSELogger): Promise<Response> {
563+
for (let attempt = 0; attempt < MAX_FETCH_RETRIES; attempt++) {
564+
try {
565+
const res = await fetch(url, init);
566+
if (res.ok || !TRANSIENT_STATUS_CODES.has(res.status) || attempt === MAX_FETCH_RETRIES - 1) {
567+
return res;
568+
}
569+
logger?.logSSEEvent(`Transient failure (${res.status}), retrying (${attempt + 1}/${MAX_FETCH_RETRIES})...`);
570+
} catch (err) {
571+
if (attempt === MAX_FETCH_RETRIES - 1) throw err;
572+
logger?.logSSEEvent(`Network error, retrying (${attempt + 1}/${MAX_FETCH_RETRIES})...`);
573+
}
574+
await new Promise(resolve => setTimeout(resolve, FETCH_RETRY_DELAY_MS));
575+
}
576+
throw new Error('fetchWithRetry: exhausted retries');
577+
}
578+
557579
/** Build the common headers for MCP bearer-token HTTP requests. */
558580
function buildMcpBearerHeaders(options: McpInvokeOptions): Record<string, string> {
559581
const headers: Record<string, string> = {
@@ -581,11 +603,7 @@ async function mcpRpcCallWithBearer(options: McpInvokeOptions, body: Record<stri
581603

582604
options.logger?.logSSEEvent(`MCP request: ${JSON.stringify(body)}`);
583605

584-
const res = await fetch(url, {
585-
method: 'POST',
586-
headers,
587-
body: JSON.stringify(body),
588-
});
606+
const res = await fetchWithRetry(url, { method: 'POST', headers, body: JSON.stringify(body) }, options.logger);
589607

590608
if (!res.ok) {
591609
const errBody = await res.text().catch(() => '');
@@ -609,11 +627,7 @@ async function mcpRpcNotifyWithBearer(options: McpInvokeOptions, body: Record<st
609627
const url = buildInvokeUrl(options.region, options.runtimeArn);
610628
const headers = buildMcpBearerHeaders(options);
611629

612-
const res = await fetch(url, {
613-
method: 'POST',
614-
headers,
615-
body: JSON.stringify(body),
616-
});
630+
const res = await fetchWithRetry(url, { method: 'POST', headers, body: JSON.stringify(body) }, options.logger);
617631

618632
if (!res.ok) {
619633
const errBody = await res.text().catch(() => '');

0 commit comments

Comments
 (0)