Add support for AWS Bedrock in MI Copilot#996
Add support for AWS Bedrock in MI Copilot#996
Conversation
|
Important Review skippedToo many files! This PR contains 296 files, which is 146 over the limit of 150. You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds AWS Bedrock authentication: types, state-machine flows, credential validation service, Bedrock connection and cache handling, UI for entering AWS credentials, and a Bedrock SDK dependency. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant UI as Login UI
participant SM as AI State Machine
participant Auth as Auth Service
participant Bedrock as AWS Bedrock
User->>UI: Click "Enter AWS Bedrock credentials"
UI->>SM: AUTH_WITH_AWS_BEDROCK
SM->>SM: enter awsBedrockFlow
User->>UI: Submit credentials
UI->>SM: SUBMIT_AWS_CREDENTIALS(payload)
SM->>SM: enter validatingAwsCredentials
SM->>Auth: validateAwsCredentials(payload)
Auth->>Bedrock: test generateText call
alt valid
Bedrock-->>Auth: success
Auth-->>SM: token / store secrets
SM->>UI: authenticated state
else invalid
Bedrock-->>Auth: error
Auth-->>SM: throw error
SM->>UI: return to awsBedrockFlow with error
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Areas requiring extra attention:
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (5)
workspaces/mi/mi-extension/src/ai-panel/copilot/connection.ts (1)
19-20: Bedrock integration looks coherent; verify region prefixes, cache shape, and consider minor refactors
getBedrockRegionalPrefixcoversus,us-gov,eu,ap, and mapsca/satous, butvalidRegionsinauth.tsalso include prefixes likeme-andaf-. It would be safer to explicitly handle those or at least confirm that defaulting them to'us'matches the Bedrock inference profile naming rules.- For the fallback branch you use
console.warn; elsewhere this module relies onlogInfo/logError. Consider reusing the shared logger for consistency.- In the AWS_BEDROCK branch of
getAnthropicClientyou recreate the Bedrock client on each call and return early instead of caching like the Anthropic client. That’s functionally fine but may be slightly inefficient; caching the factory (or at least updatingcachedAuthMethod) would bring symmetry and avoid repeated setup.getProviderCacheControlnow returns{ bedrock: { cachePoint: { type: 'default' } } }for AWS_BEDROCK and{ anthropic: { cacheControl: { type: 'ephemeral' } } }otherwise; please double‑check these shapes and keys against the latest@ai-sdk/*provider cache control docs to ensure they’re interpreted correctly.Also applies to: 32-54, 177-205, 223-232
workspaces/mi/mi-extension/src/ai-panel/aiMachine.ts (1)
25-26: AWS Bedrock auth flow mirrors API key flow; consider handlingCANCELfor parityThe new Bedrock path (
AUTH_WITH_AWS_BEDROCK→awsBedrockFlow→validatingAwsCredentialswithvalidateAwsCredentialsService) is wired consistently with the existing API key flow and uses the same error‑clearing patterns.One small consistency gap:
awsBedrockFlowhandlesCANCEL_LOGINbut notCANCEL, whereas bothssoFlowandapiKeyFlowhandle both events. IfAI_EVENT_TYPE.CANCELcan be emitted while in the Bedrock flow, adding a matching transition to#mi-ai.Unauthenticated(clearingloginMethodanderrorMessage) would keep behavior uniform.Also applies to: 140-146, 162-165, 253-288, 452-464, 480-481
workspaces/mi/mi-core/src/state-machine-types.ts (1)
193-198: AwsBedrockSecrets and AuthCredentials union are well-shaped; consider a small doc/type polish.The
AwsBedrockSecretsstructure and its inclusion inAuthCredentialsare consistent with how Bedrock credentials are validated and stored inauth.ts. Two small follow‑ups you might consider:
- Export
AwsBedrockSecrets(if you expect other modules to type against it) instead of re-declaring the shape elsewhere.- Update the
AIUserTokencomment to mention whattokenrepresents forAWS_BEDROCK(currently it only mentions MI Intel and Anthropic).Also applies to: 204-216, 218-220
workspaces/mi/mi-extension/src/ai-panel/auth.ts (2)
110-144: AWS_BEDROCK branch in getAccessToken is functionally fine but slightly confusing semantically.Returning
accessKeyIdforLoginMethod.AWS_BEDROCKkeepsgetAccessTokenusable as a generic “do we have credentials?” probe, but the namegetAccessTokenis now a bit misleading (it’s not a bearer token for Bedrock). This is acceptable, just be aware at call sites not to treat this as an HTTP Authorization token for AWS requests.
415-485: AWS credential validation flow is solid; consider relaxing heuristics and improving logging/UX.The overall approach—lightweight Bedrock call via
generateText, then persisting secrets in VS Code’s secret storage—is reasonable. A few refinements would make it more robust:
- The access key check
startsWith('AKIA') || startsWith('ASIA')is brittle; AWS uses additional valid prefixes now. A non‑empty check (and letting the Bedrock call be the source of truth) is safer.- Maintaining a hard‑coded
validRegionslist risks divergence from actual Bedrock support and from any other region list you use elsewhere. Consider centralizing this list or deferring to the provider/AWS errors for unsupported regions.- Use the existing
logErrorhelper instead ofconsole.errorso Bedrock failures end up in the same log channel and you can better control what (if anything) is emitted about the error object.- Optionally, mirror
validateApiKeyby settingmaxRetries: 0and mapping common Bedrock error codes to friendlier user messages (invalid credentials vs. region not enabled vs. network issues), to avoid generic “Validation failed” UX.None of these are blockers, but addressing them would reduce future breakage and improve the user experience.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
common/config/rush/pnpm-lock.yamlis excluded by!**/pnpm-lock.yaml
📒 Files selected for processing (9)
workspaces/mi/mi-core/src/state-machine-types.ts(5 hunks)workspaces/mi/mi-extension/package.json(1 hunks)workspaces/mi/mi-extension/src/ai-panel/aiMachine.ts(6 hunks)workspaces/mi/mi-extension/src/ai-panel/auth.ts(3 hunks)workspaces/mi/mi-extension/src/ai-panel/copilot/connection.ts(4 hunks)workspaces/mi/mi-extension/src/ai-panel/copilot/suggestions/suggestions.ts(1 hunks)workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx(3 hunks)workspaces/mi/mi-visualizer/src/views/AIPanel/index.tsx(1 hunks)workspaces/mi/mi-visualizer/src/views/LoggedOutWindow/index.tsx(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (4)
workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx (1)
workspaces/common-libs/ui-toolkit/src/components/Codicon/Codicon.tsx (1)
Codicon(44-56)
workspaces/mi/mi-extension/src/ai-panel/aiMachine.ts (2)
workspaces/mi/mi-core/src/state-machine-types.ts (1)
AIMachineContext(222-227)workspaces/mi/mi-extension/src/ai-panel/auth.ts (1)
validateAwsCredentials(415-485)
workspaces/mi/mi-extension/src/ai-panel/copilot/connection.ts (1)
workspaces/mi/mi-extension/src/ai-panel/auth.ts (2)
getAwsBedrockCredentials(149-160)getLoginMethod(101-104)
workspaces/mi/mi-extension/src/ai-panel/auth.ts (2)
workspaces/mi/mi-core/src/state-machine-types.ts (2)
AIUserToken(218-220)AuthCredentials(204-216)workspaces/mi/mi-extension/src/ai-panel/copilot/connection.ts (1)
getBedrockRegionalPrefix(37-54)
🔇 Additional comments (8)
workspaces/mi/mi-visualizer/src/views/LoggedOutWindow/index.tsx (1)
143-145: AWS Bedrock login entrypoint wiring looks correctThe new handler and button cleanly mirror the existing Anthropic API key flow and correctly dispatch
AUTH_WITH_AWS_BEDROCKviarpcClient.Also applies to: 175-175
workspaces/mi/mi-visualizer/src/views/AIPanel/index.tsx (1)
59-64: Authenticating substate routing updated correctly for AWS BedrockIncluding
validatingAwsCredentialsinisValidatingand routing bothawsBedrockFlowandvalidatingAwsCredentialstoWaitingForLoginSectionkeeps the Bedrock flow consistent with the API key flow.workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx (1)
197-205: AWS Bedrock credential form wiring and validation gating look solidState management for
awsCredentials, the visibility toggles, andconnectWithAwsCredentials(including trimming and disabling the primary button when the form is incomplete or validating) are consistent and align with the state machine’sSUBMIT_AWS_CREDENTIALSflow.Also applies to: 211-223, 264-269, 275-285, 348-460
workspaces/mi/mi-core/src/state-machine-types.ts (3)
112-119: AWS Bedrock substates are wired consistently into the AI machine state.Adding
'awsBedrockFlow'and'validatingAwsCredentials'underAuthenticatingfits the existing hierarchical pattern and keeps the state space explicit and type-safe. No issues here.
125-146: New AWS Bedrock AI events and payload typing look correct.
AUTH_WITH_AWS_BEDROCKandSUBMIT_AWS_CREDENTIALSare added to bothAI_EVENT_TYPEandAIMachineEventMapwith a concrete payload shape, which will give good type safety for the new flow. Just ensure all senders use thepayloadobject with these exact field names.Also applies to: 148-174
182-186: LoginMethod enum extension aligns with the auth flows.Introducing
ANTHROPIC_KEYandAWS_BEDROCKhere matches the usages inauth.ts(validation and token retrieval) and keeps the login methods centralized. No further changes needed in this enum.workspaces/mi/mi-extension/src/ai-panel/auth.ts (2)
30-41: New Bedrock imports and regional helper usage are consistent.
createAmazonBedrockandgetBedrockRegionalPrefixare used only in the Bedrock validation path and keep the provider‑specific logic localized. No structural issues with these additions.
146-161: getAwsBedrockCredentials helper matches stored secret shape and keeps consumers clean.The helper correctly narrows on
LoginMethod.AWS_BEDROCKand returns the exact credential fields needed by callers, which keeps Bedrock-specific logic out of generic auth code. Looks good.
| role: "system" as const, | ||
| content: systemPrompt, | ||
| providerOptions: cacheOptions, // Cache system prompt only | ||
| providerOptions: cacheOptions as any, // Cache system prompt only |
There was a problem hiding this comment.
🧩 Analysis chain
Avoid as any — properly type the provider options.
Using as any bypasses TypeScript's type safety and can hide type errors. With the AWS Bedrock integration, the proper approach is to extend the provider options type to accommodate both Anthropic and AWS Bedrock options, rather than casting to any.
Please verify the return type of getProviderCacheControl() and ensure the providerOptions field is properly typed to accept the returned value:
🏁 Script executed:
#!/bin/bash
# Description: Check the return type of getProviderCacheControl and how providerOptions is typed in the AI SDK
# Find the definition of getProviderCacheControl
ast-grep --pattern 'getProviderCacheControl($_) {
$$$
}'
# Also search for its type signature
rg -n -A 5 'getProviderCacheControl'Length of output: 16904
---
🏁 Script executed:
```shell
#!/bin/bash
# Get the full getProviderCacheControl function with all return cases
rg -n -A 20 'export const getProviderCacheControl = async' workspaces/mi/mi-extension/src/ai-panel/copilot/connection.ts
# Also check the full context around line 76 in suggestions.ts
rg -n -B 5 -A 5 'providerOptions: cacheOptions as any' workspaces/mi/mi-extension/src/ai-panel/copilot/suggestions/suggestions.ts
Length of output: 216
---
🏁 Script executed:
```shell
#!/bin/bash
# Get the full getProviderCacheControl function with all return cases
rg -n -A 20 'export const getProviderCacheControl = async' workspaces/mi/mi-extension/src/ai-panel/copilot/connection.ts
# Also check the full context around line 76 in suggestions.ts (the as any cast)
rg -n -B 10 -A 5 'providerOptions: cacheOptions as any' workspaces/mi/mi-extension/src/ai-panel/copilot/suggestions/suggestions.ts
# Also search for the type definition of ProviderCacheOptions if it exists
rg -n 'type ProviderCacheOptions|interface ProviderCacheOptions' workspaces/mi/mi-extension/src/ai-panel/copilot/
Length of output: 939
---
🏁 Script executed:
```shell
#!/bin/bash
# Check how other files use cacheOptions - see if they also use as any
rg -n 'providerOptions.*cacheOptions' workspaces/mi/mi-extension/src/ai-panel/copilot/
# Check imports to understand what types are being used
rg -n 'import.*streamText|import.*ModelMessage' workspaces/mi/mi-extension/src/ai-panel/copilot/suggestions/suggestions.ts
# Check the broader pattern - are there type definitions for provider options?
rg -n 'providerOptions' workspaces/mi/mi-extension/src/ai-panel/copilot/suggestions/suggestions.ts
Length of output: 921
Remove the unnecessary as any cast — it's inconsistent with the rest of the codebase.
The same cacheOptions from getProviderCacheControl() is used without as any in five other files: unit_test_case_generate.ts, generation/generations.ts, diagnostics/diagnostics.ts, and connectors/connectors.ts. Remove the cast on line 76 to maintain consistency and preserve type safety:
providerOptions: cacheOptions, // Cache system prompt only
🤖 Prompt for AI Agents
In workspaces/mi/mi-extension/src/ai-panel/copilot/suggestions/suggestions.ts
around line 76, remove the unnecessary `as any` type cast on `providerOptions:
cacheOptions as any` and change it to `providerOptions: cacheOptions` to match
other usages and preserve type safety; update only that line so it uses the
existing `cacheOptions` value without casting.
workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx (1)
211-223: Consider adding format validation for AWS credentials.The validation only checks for non-empty fields, unlike the Anthropic flow which validates key format and length. Consider adding client-side checks for AWS credential formats (e.g., access keys typically start with
AKIAorASIA, regions follow patterns likeus-east-1) to provide immediate feedback before backend validation.Example validation you could add:
const connectWithAwsCredentials = () => { + const trimmedAccessKey = awsCredentials.accessKeyId.trim(); + const trimmedSecretKey = awsCredentials.secretAccessKey.trim(); + const trimmedRegion = awsCredentials.region.trim(); + + if (!trimmedAccessKey || !trimmedSecretKey || !trimmedRegion) { + return; + } + + // Basic AWS access key format validation + if (!/^(AKIA|ASIA)[A-Z0-9]{16}$/.test(trimmedAccessKey)) { + // Set client error state similar to Anthropic flow + return; + } + + // Basic AWS region format validation + if (!/^[a-z]{2}-[a-z]+-\d+$/.test(trimmedRegion)) { + // Set client error state for invalid region format + return; + } + - if (awsCredentials.accessKeyId.trim() && awsCredentials.secretAccessKey.trim() && awsCredentials.region.trim()) { rpcClient.sendAIStateEvent({ type: AI_EVENT_TYPE.SUBMIT_AWS_CREDENTIALS, payload: { - accessKeyId: awsCredentials.accessKeyId.trim(), - secretAccessKey: awsCredentials.secretAccessKey.trim(), - region: awsCredentials.region.trim(), + accessKeyId: trimmedAccessKey, + secretAccessKey: trimmedSecretKey, + region: trimmedRegion, sessionToken: awsCredentials.sessionToken.trim() || undefined }, }); - } };
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx(3 hunks)
🔇 Additional comments (2)
workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx (2)
197-205: LGTM!The AWS credentials state is well-structured, follows the existing pattern for sensitive field visibility, and correctly separates concerns.
264-285: LGTM!The credential change handler correctly uses
keyoffor type safety and functional state updates. The visibility toggle handlers follow the established pattern and work correctly.
workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx
Show resolved
Hide resolved
…eAccordion and ResourceAccordionV2 for loading state management
…its usage in DiagramWrapper for resource configuration display
[BI]Add sql support for expression editor
Fix completions for method access
…ed-option Fix Boolean editor and enum editor to provide none selected option
Fix build failures in bi-extension and ballerina-graphql-design-diagram
Merge "micro-integrator-3.1.3" into "release-mi"
Fix PR comments
…atch-1 Update broken documentation link in BI extension README
Merge "micro-integrator-3.1.4" into "stable/mi"
Automatic Sync: Merge changes from stable/mi to main
Add MCP tools support
Add CodeRabbit config to review PRs on feature/*, release/*, and hotfix/* branches
Update Java version compatibility mapping for MI 4.6.0
Fix vulnerabilities
|
|
|
Closing this PR as it is a duplicate of #1499 |
18 participants
Implements: wso2/mi-vscode#1350
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.