Skip to content

Commit fd6e9a8

Browse files
authored
feat(amazonq): added mcp admin level configuration with GetProfile (#2000)
* feat(amazonq): added mcp admin level configuration with GetProfile * feat(amazonq): added UX message for mcp admin control * test: add unit tests for ProfileStatusMonitor static functionality * test: add comprehensive unit tests for MCP admin control features * fix: fix to wait for serviceManager is initialzied to initialize mcp managers * fix: fix for unit test failures * fix: fix for UI changes * fix(amazonq): fix to to rename mcp enabled function and max time limit of 10 seconds * fix: fix to add async initialization for mcp manager * fix: fix to move action buttons to mcpEventHandler for listMcpServers * fix: added try and catch block for mcp initialization * fix: fix for merge conflicts * fix: fix for test failure * fix: remove the unnecessary feature flag * fix: fix for mynah test failure * fix: fix to retry function to common util * fix: fix to add retryUtils
1 parent 817cfe2 commit fd6e9a8

File tree

14 files changed

+899
-38
lines changed

14 files changed

+899
-38
lines changed

chat-client/src/client/mcpMynahUi.test.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,8 @@ describe('McpMynahUi', () => {
107107
assert.strictEqual(callArgs.detailedList.header.description, 'Test Description')
108108
assert.deepStrictEqual(callArgs.detailedList.header.status, { status: 'success' })
109109

110-
// Verify the actions in the header
111-
assert.strictEqual(callArgs.detailedList.header.actions.length, 2)
112-
assert.strictEqual(callArgs.detailedList.header.actions[0].id, 'add-new-mcp')
113-
assert.strictEqual(callArgs.detailedList.header.actions[1].id, 'refresh-mcp-list')
110+
// Verify the actions in the header (no default actions are added when header is provided)
111+
assert.strictEqual(callArgs.detailedList.header.actions.length, 0)
114112

115113
// Verify the list structure
116114
assert.strictEqual(callArgs.detailedList.list.length, 1)

chat-client/src/client/mcpMynahUi.ts

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -272,20 +272,11 @@ export class McpMynahUi {
272272
title: params.header.title,
273273
description: params.header.description,
274274
status: params.header.status,
275-
actions: [
276-
{
277-
id: MCP_IDS.ADD_NEW,
278-
icon: toMynahIcon('plus'),
279-
status: 'clear',
280-
description: 'Add new MCP',
281-
},
282-
{
283-
id: MCP_IDS.REFRESH_LIST,
284-
icon: toMynahIcon('refresh'),
285-
status: 'clear',
286-
description: 'Refresh MCP servers',
287-
},
288-
],
275+
actions:
276+
params.header.actions?.map(action => ({
277+
...action,
278+
icon: action.icon ? toMynahIcon(action.icon) : undefined,
279+
})) || [],
289280
}
290281
: undefined,
291282
filterOptions: params.filterOptions?.map(filter => ({

core/aws-lsp-core/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,3 +19,4 @@ export * as workspaceUtils from './util/workspaceUtils'
1919
export * as processUtils from './util/processUtils'
2020
export * as collectionUtils from './util/collectionUtils'
2121
export * as loggingUtils from './util/loggingUtils'
22+
export * as retryUtils from './util/retryUtils'
Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates.
3+
* All Rights Reserved. SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import { expect } from 'chai'
7+
import * as sinon from 'sinon'
8+
import { retryWithBackoff, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY } from './retryUtils'
9+
10+
describe('retryUtils', () => {
11+
let clock: sinon.SinonFakeTimers
12+
13+
beforeEach(() => {
14+
clock = sinon.useFakeTimers()
15+
})
16+
17+
afterEach(() => {
18+
clock.restore()
19+
})
20+
21+
describe('retryWithBackoff', () => {
22+
it('should return result on first success', async () => {
23+
const fn = sinon.stub().resolves('success')
24+
25+
const result = await retryWithBackoff(fn)
26+
27+
expect(result).to.equal('success')
28+
expect(fn.callCount).to.equal(1)
29+
})
30+
31+
it('should retry on retryable errors', async () => {
32+
const fn = sinon.stub()
33+
fn.onFirstCall().rejects({ code: 'ThrottlingException' })
34+
fn.onSecondCall().resolves('success')
35+
36+
const promise = retryWithBackoff(fn)
37+
await clock.tickAsync(DEFAULT_BASE_DELAY)
38+
const result = await promise
39+
40+
expect(result).to.equal('success')
41+
expect(fn.callCount).to.equal(2)
42+
})
43+
44+
it('should not retry on non-retryable client errors', async () => {
45+
const error = { statusCode: 404 }
46+
const fn = sinon.stub().rejects(error)
47+
48+
try {
49+
await retryWithBackoff(fn)
50+
expect.fail('Expected function to throw')
51+
} catch (e) {
52+
expect(e).to.equal(error)
53+
}
54+
expect(fn.callCount).to.equal(1)
55+
})
56+
57+
it('should retry on server errors', async () => {
58+
const fn = sinon.stub()
59+
fn.onFirstCall().rejects({ statusCode: 500 })
60+
fn.onSecondCall().resolves('success')
61+
62+
const promise = retryWithBackoff(fn)
63+
await clock.tickAsync(DEFAULT_BASE_DELAY)
64+
const result = await promise
65+
66+
expect(result).to.equal('success')
67+
expect(fn.callCount).to.equal(2)
68+
})
69+
70+
it('should use exponential backoff by default', async () => {
71+
const fn = sinon.stub()
72+
const error = { code: 'ThrottlingException' }
73+
fn.onFirstCall().rejects(error)
74+
fn.onSecondCall().rejects(error)
75+
76+
const promise = retryWithBackoff(fn)
77+
78+
// First retry after baseDelay * 1
79+
await clock.tickAsync(DEFAULT_BASE_DELAY)
80+
// Second retry after baseDelay * 2
81+
await clock.tickAsync(DEFAULT_BASE_DELAY * 2)
82+
83+
try {
84+
await promise
85+
expect.fail('Expected function to throw')
86+
} catch (e) {
87+
expect(e).to.equal(error)
88+
}
89+
expect(fn.callCount).to.equal(DEFAULT_MAX_RETRIES)
90+
})
91+
92+
it('should respect custom maxRetries', async () => {
93+
const error = { code: 'ThrottlingException' }
94+
const fn = sinon.stub().rejects(error)
95+
96+
try {
97+
await retryWithBackoff(fn, { maxRetries: 1 })
98+
expect.fail('Expected function to throw')
99+
} catch (e) {
100+
expect(e).to.equal(error)
101+
}
102+
expect(fn.callCount).to.equal(1)
103+
})
104+
105+
it('should use custom isRetryable function', async () => {
106+
const error = { custom: 'error' }
107+
const fn = sinon.stub().rejects(error)
108+
const isRetryable = sinon.stub().returns(false)
109+
110+
try {
111+
await retryWithBackoff(fn, { isRetryable })
112+
expect.fail('Expected function to throw')
113+
} catch (e) {
114+
expect(e).to.equal(error)
115+
}
116+
expect(fn.callCount).to.equal(1)
117+
expect(isRetryable.calledWith(error)).to.equal(true)
118+
})
119+
})
120+
})
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates.
3+
* All Rights Reserved. SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
// Default retry configuration constants
7+
export const DEFAULT_MAX_RETRIES = 2
8+
export const DEFAULT_BASE_DELAY = 500
9+
export const DEFAULT_EXPONENTIAL_BACKOFF = true
10+
11+
// HTTP status code constants
12+
const CLIENT_ERROR_MIN = 400
13+
const CLIENT_ERROR_MAX = 500
14+
const INTERNAL_SERVER_ERROR = 500
15+
const SERVICE_UNAVAILABLE = 503
16+
17+
// AWS error code constants
18+
const THROTTLING_EXCEPTION = 'ThrottlingException'
19+
const INTERNAL_SERVER_EXCEPTION = 'InternalServerException'
20+
21+
export interface RetryOptions {
22+
/** Maximum number of retry attempts (default: DEFAULT_MAX_RETRIES) */
23+
maxRetries?: number
24+
/** Base delay in milliseconds (default: DEFAULT_BASE_DELAY) */
25+
baseDelay?: number
26+
/** Whether to use exponential backoff (default: DEFAULT_EXPONENTIAL_BACKOFF) */
27+
exponentialBackoff?: boolean
28+
/** Custom function to determine if an error is retryable */
29+
isRetryable?: (error: any) => boolean
30+
}
31+
32+
/**
33+
* Default AWS error retry logic
34+
*/
35+
function defaultIsRetryable(error: any): boolean {
36+
const errorCode = error.code || error.name
37+
const statusCode = error.statusCode
38+
39+
// Fast fail on non-retryable client errors (except throttling)
40+
if (statusCode >= CLIENT_ERROR_MIN && statusCode < CLIENT_ERROR_MAX && errorCode !== THROTTLING_EXCEPTION) {
41+
return false
42+
}
43+
44+
// Retry on throttling, server errors, and specific status codes
45+
return (
46+
errorCode === THROTTLING_EXCEPTION ||
47+
errorCode === INTERNAL_SERVER_EXCEPTION ||
48+
statusCode === INTERNAL_SERVER_ERROR ||
49+
statusCode === SERVICE_UNAVAILABLE
50+
)
51+
}
52+
53+
/**
54+
* Executes a function with retry logic and exponential backoff
55+
*/
56+
export async function retryWithBackoff<T>(fn: () => Promise<T>, options: RetryOptions = {}): Promise<T> {
57+
const {
58+
maxRetries = DEFAULT_MAX_RETRIES,
59+
baseDelay = DEFAULT_BASE_DELAY,
60+
exponentialBackoff = DEFAULT_EXPONENTIAL_BACKOFF,
61+
isRetryable = defaultIsRetryable,
62+
} = options
63+
64+
for (let attempt = 0; attempt < maxRetries; attempt++) {
65+
try {
66+
return await fn()
67+
} catch (error: any) {
68+
if (!isRetryable(error) || attempt === maxRetries - 1) {
69+
throw error
70+
}
71+
72+
const delay = exponentialBackoff ? baseDelay * (attempt + 1) : baseDelay
73+
await new Promise(resolve => setTimeout(resolve, delay))
74+
}
75+
}
76+
throw new Error('Retry failed')
77+
}

server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -516,6 +516,37 @@
516516
],
517517
"documentation": "<p>API to get code transformation status.</p>"
518518
},
519+
"GetProfile": {
520+
"name": "GetProfile",
521+
"http": {
522+
"method": "POST",
523+
"requestUri": "/"
524+
},
525+
"input": {
526+
"shape": "GetProfileRequest"
527+
},
528+
"output": {
529+
"shape": "GetProfileResponse"
530+
},
531+
"errors": [
532+
{
533+
"shape": "ThrottlingException"
534+
},
535+
{
536+
"shape": "ResourceNotFoundException"
537+
},
538+
{
539+
"shape": "InternalServerException"
540+
},
541+
{
542+
"shape": "ValidationException"
543+
},
544+
{
545+
"shape": "AccessDeniedException"
546+
}
547+
],
548+
"documentation": "<p>Get the requested CodeWhisperer profile.</p>"
549+
},
519550
"GetUsageLimits": {
520551
"name": "GetUsageLimits",
521552
"http": {
@@ -3126,6 +3157,24 @@
31263157
}
31273158
}
31283159
},
3160+
"GetProfileRequest": {
3161+
"type": "structure",
3162+
"required": ["profileArn"],
3163+
"members": {
3164+
"profileArn": {
3165+
"shape": "ProfileArn"
3166+
}
3167+
}
3168+
},
3169+
"GetProfileResponse": {
3170+
"type": "structure",
3171+
"required": ["profile"],
3172+
"members": {
3173+
"profile": {
3174+
"shape": "ProfileInfo"
3175+
}
3176+
}
3177+
},
31293178
"GetTaskAssistCodeGenerationRequest": {
31303179
"type": "structure",
31313180
"required": ["conversationId", "codeGenerationId"],
@@ -3787,6 +3836,15 @@
37873836
"type": "long",
37883837
"box": true
37893838
},
3839+
"MCPConfiguration": {
3840+
"type": "structure",
3841+
"required": ["toggle"],
3842+
"members": {
3843+
"toggle": {
3844+
"shape": "OptInFeatureToggle"
3845+
}
3846+
}
3847+
},
37903848
"MemoryEntry": {
37913849
"type": "structure",
37923850
"required": ["id", "memoryEntryString", "metadata"],
@@ -3995,6 +4053,9 @@
39954053
},
39964054
"workspaceContext": {
39974055
"shape": "WorkspaceContext"
4056+
},
4057+
"mcpConfiguration": {
4058+
"shape": "MCPConfiguration"
39984059
}
39994060
}
40004061
},
@@ -4189,6 +4250,30 @@
41894250
}
41904251
}
41914252
},
4253+
"ProfileInfo": {
4254+
"type": "structure",
4255+
"required": ["arn"],
4256+
"members": {
4257+
"arn": {
4258+
"shape": "ProfileArn"
4259+
},
4260+
"profileName": {
4261+
"shape": "ProfileName"
4262+
},
4263+
"description": {
4264+
"shape": "ProfileDescription"
4265+
},
4266+
"status": {
4267+
"shape": "ProfileStatus"
4268+
},
4269+
"profileType": {
4270+
"shape": "ProfileType"
4271+
},
4272+
"optInFeatures": {
4273+
"shape": "OptInFeatures"
4274+
}
4275+
}
4276+
},
41924277
"ProfileArn": {
41934278
"type": "string",
41944279
"max": 950,

0 commit comments

Comments
 (0)