Skip to content

Add support for AWS Bedrock in MI Copilot #996

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ravindu25 wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Add support for AWS Bedrock in MI Copilot #996

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ef7fd32
Add support for AWS Bedrock in MI Copilot
ravindu25 Nov 24, 2025
9294e39
Merge branch 'main' into main
ravindu25 Nov 24, 2025
986aaca
add pnpm lock file
ravindu25 Nov 24, 2025
123650c
Merge branch 'main' of https://github.com/ravindu25/vscode-extensions…
ravindu25 Nov 24, 2025
2feaba5
Fix typo
ravindu25 Nov 25, 2025
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
13 changes: 13 additions & 0 deletions common/config/rush/pnpm-lock.yaml
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

25 changes: 23 additions & 2 deletions workspaces/mi/mi-core/src/state-machine-types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ export type MachineStateValue =
export type AIMachineStateValue =
| 'Initialize' // (checking auth, first load)
| 'Unauthenticated' // (show login window)
| { Authenticating: 'determineFlow' | 'ssoFlow' | 'apiKeyFlow' | 'validatingApiKey' } // hierarchical substates
| { Authenticating: 'determineFlow' | 'ssoFlow' | 'apiKeyFlow' | 'validatingApiKey' | 'awsBedrockFlow' | 'validatingAwsCredentials' } // hierarchical substates
| 'Authenticated' // (ready, main view)
| 'UsageExceeded' // (free usage quota exceeded, prompt user to set API key)
| 'Disabled' // (optional: if AI Chat is globally unavailable)
Expand All @@ -127,6 +127,8 @@ export enum AI_EVENT_TYPE {
LOGIN = "LOGIN",
AUTH_WITH_API_KEY = 'AUTH_WITH_API_KEY',
SUBMIT_API_KEY = 'SUBMIT_API_KEY',
AUTH_WITH_AWS_BEDROCK = 'AUTH_WITH_AWS_BEDROCK',
SUBMIT_AWS_CREDENTIALS = 'SUBMIT_AWS_CREDENTIALS',
SIGN_IN_SUCCESS = "SIGN_IN_SUCCESS",
LOGOUT = "LOGOUT",
SILENT_LOGOUT = "SILENT_LOGOUT",
Expand All @@ -148,6 +150,13 @@ export type AIMachineEventMap = {
[AI_EVENT_TYPE.LOGIN]: undefined;
[AI_EVENT_TYPE.AUTH_WITH_API_KEY]: undefined;
[AI_EVENT_TYPE.SUBMIT_API_KEY]: { apiKey: string };
[AI_EVENT_TYPE.AUTH_WITH_AWS_BEDROCK]: undefined;
[AI_EVENT_TYPE.SUBMIT_AWS_CREDENTIALS]: {
accessKeyId: string;
secretAccessKey: string;
region: string;
sessionToken?: string;
};
[AI_EVENT_TYPE.SIGN_IN_SUCCESS]: undefined;
[AI_EVENT_TYPE.LOGOUT]: undefined;
[AI_EVENT_TYPE.SILENT_LOGOUT]: undefined;
Expand All @@ -172,14 +181,22 @@ export type AIMachineSendableEvent =

export enum LoginMethod {
MI_INTEL = 'miIntel',
ANTHROPIC_KEY = 'anthropic_key'
ANTHROPIC_KEY = 'anthropic_key',
AWS_BEDROCK = 'aws_bedrock'
}

interface MIIntelSecrets {
accessToken: string;
refreshToken: string;
}

interface AwsBedrockSecrets {
accessKeyId: string;
secretAccessKey: string;
region: string;
sessionToken?: string;
}

interface AnthropicKeySecrets {
apiKey: string;
}
Expand All @@ -192,6 +209,10 @@ export type AuthCredentials =
| {
loginMethod: LoginMethod.ANTHROPIC_KEY;
secrets: AnthropicKeySecrets;
}
| {
loginMethod: LoginMethod.AWS_BEDROCK;
secrets: AwsBedrockSecrets;
};

export interface AIUserToken {
Expand Down
3 changes: 2 additions & 1 deletion workspaces/mi/mi-extension/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1060,6 +1060,7 @@
"vscode-messenger-common": "^0.5.1",
"xml2js": "0.6.2",
"xstate": "^4.38.3",
"zod": "^3.24.1"
"zod": "^3.24.1",
"@ai-sdk/amazon-bedrock": "~3.0.56"
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

Verify version pinning consistency and package availability.

The new @ai-sdk/amazon-bedrock dependency uses a tilde (~3.0.56) for version pinning, while the existing @ai-sdk/anthropic dependency uses a caret (^2.0.35). This inconsistency may be intentional, but should be clarified. Additionally, verify that the specified version exists on npm and contains no known vulnerabilities.

🏁 Script executed:

#!/bin/bash
# Description: Verify the AWS Bedrock SDK version exists and check for security advisories

# Check if the package version exists on npm
npm view @ai-sdk/[email protected] version 2>/dev/null || echo "Version 3.0.56 not found"

# Get latest version for comparison
echo "Latest version available:"
npm view @ai-sdk/amazon-bedrock@latest version 2>/dev/null

# Check for known security advisories (if npm audit is available)
echo "Checking for security issues..."
npm audit --package-lock-only @ai-sdk/[email protected] 2>/dev/null || echo "Audit not available in this context"

Length of output: 95

🌐 Web query:

@ai-sdk/amazon-bedrock compatibility peerDependencies ai-sdk version requirements

💡 Result:

Short answer: @ai-sdk/amazon-bedrock does not list the core "ai" package as a peerDependency — instead it depends on the provider framework packages (e.g. @ai-sdk/provider, @ai-sdk/provider-utils) and typically requires the matching family of AI SDK packages for the same major release. Check the exact package.json for the version you plan to install, but common facts:

  • The Amazon Bedrock provider package is published separately (e.g. versions 2.x / 3.x / 3.0.16 / 3.0.33 shown on npm) and is installed alongside the core "ai" package. [1][2]
  • The provider packages declare dependencies on @ai-sdk/provider and @ai-sdk/provider-utils (not a direct peerDependency on "ai"); some builds only declare an optional peerDependency on zod. See the provider docs and package metadata. [1][3]
  • Recommendation: match provider + core major versions (use the same AI SDK major series — e.g. v4 vs v5) and inspect the specific @ai-sdk/amazon-bedrock package.json for the peerDependencies for the exact published version you’ll install. You can view metadata on the package page or the docs. [1][2][3]

Sources:
[1] AI SDK — Amazon Bedrock provider docs.
[2] @ai-sdk/amazon-bedrock npm listing (package + version info).
[3] package metadata / registry snapshots showing dependencies (e.g. provider / provider-utils and zod peerDependency).

🏁 Script executed:

#!/bin/bash
# Fetch detailed package metadata for @ai-sdk/[email protected]
npm view @ai-sdk/[email protected] --json | jq '.peerDependencies, .dependencies' 2>/dev/null || echo "Failed to fetch metadata"

Length of output: 238

🌐 Web query:

ai-sdk major version 5 amazon bedrock provider 3 compatibility

💡 Result:

Short answer: Use the provider version set required by AI SDK v5 — @ai-sdk/amazon-bedrock v2.0.x. AI SDK v5 expects @ai-sdk/* packages at 2.x (and [email protected]), so an @ai-sdk/amazon-bedrock v3 release is not the documented target for AI SDK v5 and may be incompatible. [1][2]

References

  • AI SDK 4→5 migration (lists required package versions for v5, including @ai-sdk/* = 2.0.0). [1]
  • Amazon Bedrock provider docs for the v5-era SDK (shows @ai-sdk/amazon-bedrock usage/migration guidance). [2]

Update @ai-sdk/amazon-bedrock to version 2.x for AI SDK v5 compatibility.

The specified version @ai-sdk/[email protected] is incompatible with the core ai@^5.0.76 dependency. AI SDK v5 expects @ai-sdk/* packages at 2.x. Change the dependency to @ai-sdk/amazon-bedrock@^2.0.x to match the project's AI SDK version series.

🤖 Prompt for AI Agents 
In workspaces/mi/mi-extension/package.json around line 1064, the dependency
"@ai-sdk/amazon-bedrock": "~3.0.56" is incompatible with the project's AI SDK v5
which requires @ai-sdk packages in the 2.x series; update this dependency to use
the 2.x range (for example "@ai-sdk/amazon-bedrock": "^2.0.x") so it aligns with
ai@^5.0.76 and reinstall/update lockfile to ensure consistent dependency
resolution.

}
}
61 changes: 60 additions & 1 deletion workspaces/mi/mi-extension/src/ai-panel/aiMachine.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import * as vscode from 'vscode';
import { AIMachineStateValue, AIMachineContext, AI_EVENT_TYPE, AIMachineSendableEvent, LoginMethod } from '@wso2/mi-core';
import { AiPanelWebview } from './webview';
import { extension } from '../MIExtensionContext';
import { getAccessToken, getLoginMethod, checkToken, initiateInbuiltAuth, logout, validateApiKey } from './auth';
import { getAccessToken, getLoginMethod, checkToken, initiateInbuiltAuth, logout, validateApiKey, validateAwsCredentials } from './auth';
import { PromptObject } from '@wso2/mi-core';

export const openAIWebview = (initialPrompt?: PromptObject) => {
Expand Down Expand Up @@ -137,6 +137,12 @@ const aiMachine = createMachine<AIMachineContext, AIMachineSendableEvent>({
actions: assign({
loginMethod: (_ctx) => LoginMethod.ANTHROPIC_KEY
})
},
[AI_EVENT_TYPE.AUTH_WITH_AWS_BEDROCK]: {
target: 'Authenticating',
actions: assign({
loginMethod: (_ctx) => LoginMethod.AWS_BEDROCK
})
}
}
},
Expand All @@ -153,6 +159,10 @@ const aiMachine = createMachine<AIMachineContext, AIMachineSendableEvent>({
cond: (context) => context.loginMethod === LoginMethod.ANTHROPIC_KEY,
target: 'apiKeyFlow'
},
{
cond: (context) => context.loginMethod === LoginMethod.AWS_BEDROCK,
target: 'awsBedrockFlow'
},
{
target: 'ssoFlow' // default
}
Expand Down Expand Up @@ -240,6 +250,41 @@ const aiMachine = createMachine<AIMachineContext, AIMachineSendableEvent>({
})
}
}
},
awsBedrockFlow: {
on: {
[AI_EVENT_TYPE.SUBMIT_AWS_CREDENTIALS]: {
target: 'validatingAwsCredentials',
actions: assign({
errorMessage: (_ctx) => undefined
})
},
[AI_EVENT_TYPE.CANCEL_LOGIN]: {
target: '#mi-ai.Unauthenticated',
actions: assign({
loginMethod: (_ctx) => undefined,
errorMessage: (_ctx) => undefined,
})
}
}
},
validatingAwsCredentials: {
invoke: {
id: 'validateAwsCredentials',
src: 'validateAwsCredentials',
onDone: {
target: '#mi-ai.Authenticated',
actions: assign({
errorMessage: (_ctx) => undefined,
})
},
onError: {
target: 'awsBedrockFlow',
actions: assign({
errorMessage: (_ctx, event) => event.data?.message || 'AWS credentials validation failed'
})
}
}
}
}
},
Expand Down Expand Up @@ -404,6 +449,19 @@ const validateApiKeyService = async (_context: AIMachineContext, event: any) =>
return await validateApiKey(apiKey, LoginMethod.ANTHROPIC_KEY);
};

const validateAwsCredentialsService = async (_context: AIMachineContext, event: any) => {
const { accessKeyId, secretAccessKey, region, sessionToken } = event.payload || {};
if (!accessKeyId || !secretAccessKey || !region) {
throw new Error('AWS access key ID, secret access key, and region are required');
}
return await validateAwsCredentials({
accessKeyId,
secretAccessKey,
region,
sessionToken
});
};

const getTokenAndLoginMethod = async () => {
const result = await getAccessToken();
const loginMethod = await getLoginMethod();
Expand All @@ -419,6 +477,7 @@ const aiStateService = interpret(aiMachine.withConfig({
checkWorkspaceAndToken: checkWorkspaceAndToken,
openLogin: openLogin,
validateApiKey: validateApiKeyService,
validateAwsCredentials: validateAwsCredentialsService,
getTokenAndLoginMethod: getTokenAndLoginMethod,
},
actions: {
Expand Down
92 changes: 92 additions & 0 deletions workspaces/mi/mi-extension/src/ai-panel/auth.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,10 @@ import { extension } from '../MIExtensionContext';
import * as vscode from 'vscode';
import { jwtDecode, JwtPayload } from 'jwt-decode';
import { createAnthropic } from '@ai-sdk/anthropic';
import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock';
import { generateText } from 'ai';
import { logInfo, logWarn, logError } from './copilot/logger';
import { getBedrockRegionalPrefix } from './copilot/connection';

// ==================================
// Configuration
Expand Down Expand Up @@ -132,13 +134,31 @@ export const getAccessToken = async (): Promise<string | undefined> => {
throw err;
}
}
case LoginMethod.AWS_BEDROCK:
return credentials.secrets.accessKeyId;
case LoginMethod.ANTHROPIC_KEY:
return credentials.secrets.apiKey;
}

return undefined;
};

/**
* Get AWS Bedrock credentials (accessKeyId, secretAccessKey, region, optional sessionToken)
*/
export const getAwsBedrockCredentials = async (): Promise<{
accessKeyId: string;
secretAccessKey: string;
region: string;
sessionToken?: string;
} | undefined> => {
const credentials = await getAuthCredentials();
if (!credentials || credentials.loginMethod !== LoginMethod.AWS_BEDROCK) {
return undefined;
}
return credentials.secrets;
};

/**
* Refresh MI_INTEL access token using refresh token
*/
Expand Down Expand Up @@ -392,6 +412,78 @@ export const logout = async (isUserLogout: boolean = true): Promise<void> => {
await clearAuthCredentials();
};

export const validateAwsCredentials = async (credentials: {
accessKeyId: string;
secretAccessKey: string;
region: string;
sessionToken?: string;
}): Promise<AIUserToken> => {
const { accessKeyId, secretAccessKey, region, sessionToken } = credentials;

if (!accessKeyId || !secretAccessKey || !region) {
throw new Error('AWS access key ID, secret access key, and region are required.');
}

if (!accessKeyId.startsWith('AKIA') && !accessKeyId.startsWith('ASIA')) {
throw new Error('Please enter a valid AWS access key ID.');
}

if (secretAccessKey.length < 20) {
throw new Error('Please enter a valid AWS secret access key.');
}

// List of valid AWS regions
const validRegions = [
'us-east-1', 'us-west-2', 'us-west-1', 'eu-west-1', 'eu-central-1',
'ap-southeast-1', 'ap-southeast-2', 'ap-northeast-1', 'ap-northeast-2',
'ap-south-1', 'ca-central-1', 'sa-east-1', 'eu-west-2', 'eu-west-3',
'eu-north-1', 'ap-east-1', 'me-south-1', 'af-south-1', 'ap-southeast-3'
];

if (!validRegions.includes(region)) {
throw new Error(`Invalid AWS region. Please select a valid region like us-east-1, us-west-2, etc.`);
}

try {
const bedrock = createAmazonBedrock({
region: region,
accessKeyId: accessKeyId,
secretAccessKey: secretAccessKey,
sessionToken: sessionToken,
});

// Get regional prefix based on AWS region and construct model ID
const regionalPrefix = getBedrockRegionalPrefix(region);
const modelId = `${regionalPrefix}.anthropic.claude-3-5-haiku-20241022-v1:0`;
const bedrockClient = bedrock(modelId);

// Make a minimal test call to validate credentials
await generateText({
model: bedrockClient,
maxOutputTokens: 1,
messages: [{ role: 'user', content: 'Hi' }]
});

// Store credentials
const authCredentials: AuthCredentials = {
loginMethod: LoginMethod.AWS_BEDROCK,
secrets: {
accessKeyId,
secretAccessKey,
region,
sessionToken
}
};
await storeAuthCredentials(authCredentials);

return { token: accessKeyId };

} catch (error) {
console.error('AWS Bedrock validation failed:', error);
throw new Error('Validation failed. Please check the log for more details.');
}
};

// ==================================
// Deprecated/Legacy Functions
// ==================================
Expand Down
Loading
Loading