From 7fd09f3981555440831007636fc826781cd167ae Mon Sep 17 00:00:00 2001 From: Ravindu Wegiriya Date: Wed, 25 Feb 2026 12:01:04 +0530 Subject: [PATCH 1/3] Add AWS Bedrock Login --- common/config/rush/pnpm-lock.yaml | 30 +++- .../mi/mi-core/src/state-machine-types.ts | 25 ++- workspaces/mi/mi-extension/package.json | 1 + .../mi-extension/src/ai-features/aiMachine.ts | 82 +++++++++ .../mi/mi-extension/src/ai-features/auth.ts | 97 ++++++++++- .../src/ai-features/connection.ts | 82 ++++++++- .../rpc-managers/ai-features/rpc-manager.ts | 4 +- .../views/AIPanel/component/AIChatHeader.tsx | 9 +- .../component/WaitingForLoginSection.tsx | 159 ++++++++++++++++++ .../mi-visualizer/src/views/AIPanel/index.tsx | 10 +- .../src/views/LoggedOutWindow/index.tsx | 5 + 11 files changed, 491 insertions(+), 13 deletions(-) diff --git a/common/config/rush/pnpm-lock.yaml b/common/config/rush/pnpm-lock.yaml index fc8c007d011..221bcf1ad15 100644 --- a/common/config/rush/pnpm-lock.yaml +++ b/common/config/rush/pnpm-lock.yaml @@ -3742,6 +3742,9 @@ importers: ../../workspaces/mi/mi-extension: dependencies: + '@ai-sdk/amazon-bedrock': + specifier: 4.0.4 + version: 4.0.4(zod@4.1.11) '@ai-sdk/anthropic': specifier: 3.0.46 version: 3.0.46(zod@4.1.11) @@ -25627,6 +25630,16 @@ snapshots: '@adobe/css-tools@4.4.4': {} + '@ai-sdk/amazon-bedrock@4.0.4(zod@4.1.11)': + dependencies: + '@ai-sdk/anthropic': 3.0.2(zod@4.1.11) + '@ai-sdk/provider': 3.0.1 + '@ai-sdk/provider-utils': 4.0.2(zod@4.1.11) + '@smithy/eventstream-codec': 4.2.8 + '@smithy/util-utf8': 4.2.0 + aws4fetch: 1.0.20 + zod: 4.1.11 + '@ai-sdk/amazon-bedrock@4.0.4(zod@4.1.8)': dependencies: '@ai-sdk/anthropic': 3.0.2(zod@4.1.8) @@ -25637,6 +25650,12 @@ snapshots: aws4fetch: 1.0.20 zod: 4.1.8 + '@ai-sdk/anthropic@3.0.2(zod@4.1.11)': + dependencies: + '@ai-sdk/provider': 3.0.1 + '@ai-sdk/provider-utils': 4.0.2(zod@4.1.11) + zod: 4.1.11 + '@ai-sdk/anthropic@3.0.2(zod@4.1.8)': dependencies: '@ai-sdk/provider': 3.0.1 @@ -25676,6 +25695,13 @@ snapshots: eventsource-parser: 3.0.6 zod: 4.1.11 + '@ai-sdk/provider-utils@4.0.2(zod@4.1.11)': + dependencies: + '@ai-sdk/provider': 3.0.1 + '@standard-schema/spec': 1.1.0 + eventsource-parser: 3.0.6 + zod: 4.1.11 + '@ai-sdk/provider-utils@4.0.2(zod@4.1.8)': dependencies: '@ai-sdk/provider': 3.0.1 @@ -25732,7 +25758,7 @@ snapshots: '@aws-crypto/crc32c@5.2.0': dependencies: '@aws-crypto/util': 5.2.0 - '@aws-sdk/types': 3.804.0 + '@aws-sdk/types': 3.973.1 tslib: 2.8.1 '@aws-crypto/sha1-browser@5.2.0': @@ -25766,7 +25792,7 @@ snapshots: '@aws-crypto/util@5.2.0': dependencies: - '@aws-sdk/types': 3.804.0 + '@aws-sdk/types': 3.973.1 '@smithy/util-utf8': 2.3.0 tslib: 2.8.1 diff --git a/workspaces/mi/mi-core/src/state-machine-types.ts b/workspaces/mi/mi-core/src/state-machine-types.ts index bd037113df0..990126e4a8c 100644 --- a/workspaces/mi/mi-core/src/state-machine-types.ts +++ b/workspaces/mi/mi-core/src/state-machine-types.ts @@ -113,7 +113,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) @@ -128,6 +128,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", @@ -149,6 +151,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; @@ -173,7 +182,8 @@ export type AIMachineSendableEvent = export enum LoginMethod { MI_INTEL = 'miIntel', - ANTHROPIC_KEY = 'anthropic_key' + ANTHROPIC_KEY = 'anthropic_key', + AWS_BEDROCK = 'aws_bedrock' } interface MIIntelSecrets { @@ -186,6 +196,13 @@ interface AnthropicKeySecrets { apiKey: string; } +export interface AwsBedrockSecrets { + accessKeyId: string; + secretAccessKey: string; + region: string; + sessionToken?: string; +} + export type AuthCredentials = | { loginMethod: LoginMethod.MI_INTEL; @@ -194,6 +211,10 @@ export type AuthCredentials = | { loginMethod: LoginMethod.ANTHROPIC_KEY; secrets: AnthropicKeySecrets; + } + | { + loginMethod: LoginMethod.AWS_BEDROCK; + secrets: AwsBedrockSecrets; }; export interface AIUserToken { diff --git a/workspaces/mi/mi-extension/package.json b/workspaces/mi/mi-extension/package.json index d5d5243e9d6..ef7ffb5c426 100644 --- a/workspaces/mi/mi-extension/package.json +++ b/workspaces/mi/mi-extension/package.json @@ -1069,6 +1069,7 @@ "yaml": "2.8.0" }, "dependencies": { + "@ai-sdk/amazon-bedrock": "4.0.4", "@ai-sdk/anthropic": "3.0.46", "@anthropic-ai/tokenizer": "0.0.4", "@apidevtools/json-schema-ref-parser": "12.0.2", diff --git a/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts b/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts index 540c06fbd3a..cba5bb1c676 100644 --- a/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts +++ b/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts @@ -27,6 +27,7 @@ import { initiateInbuiltAuth, logout, validateApiKey, + validateAwsCredentials, isPlatformExtensionAvailable, isDevantUserLoggedIn, getPlatformStsToken, @@ -150,6 +151,12 @@ const aiMachine = createMachine({ actions: assign({ loginMethod: (_ctx) => LoginMethod.ANTHROPIC_KEY }) + }, + [AI_EVENT_TYPE.AUTH_WITH_AWS_BEDROCK]: { + target: 'Authenticating', + actions: assign({ + loginMethod: (_ctx) => LoginMethod.AWS_BEDROCK + }) } } }, @@ -166,6 +173,10 @@ const aiMachine = createMachine({ cond: (context) => context.loginMethod === LoginMethod.ANTHROPIC_KEY, target: 'apiKeyFlow' }, + { + cond: (context) => context.loginMethod === LoginMethod.AWS_BEDROCK, + target: 'awsBedrockFlow' + }, { target: 'ssoFlow' // default } @@ -253,6 +264,48 @@ const aiMachine = createMachine({ }) } } + }, + 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, + }) + }, + [AI_EVENT_TYPE.CANCEL]: { + 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' + }) + } + } } } }, @@ -324,6 +377,13 @@ const aiMachine = createMachine({ errorMessage: (_ctx) => undefined, }) }, + [AI_EVENT_TYPE.AUTH_WITH_AWS_BEDROCK]: { + target: 'Authenticating', + actions: assign({ + loginMethod: (_ctx) => LoginMethod.AWS_BEDROCK, + errorMessage: (_ctx) => undefined, + }) + }, [AI_EVENT_TYPE.USAGE_RESET]: { target: 'Authenticated', actions: assign({ @@ -409,6 +469,11 @@ const checkWorkspaceAndToken = async (): Promise<{ workspaceSupported: boolean; if (apiKey) { tokenData = { token: apiKey, loginMethod: LoginMethod.ANTHROPIC_KEY }; } + } else if (credentials?.loginMethod === LoginMethod.AWS_BEDROCK) { + const accessKeyId = (credentials.secrets as { accessKeyId?: string })?.accessKeyId; + if (accessKeyId) { + tokenData = { token: accessKeyId, loginMethod: LoginMethod.AWS_BEDROCK }; + } } return { workspaceSupported: true, tokenData }; @@ -450,6 +515,14 @@ 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 credentials = await getAuthCredentials(); if (!credentials) { @@ -472,6 +545,14 @@ const getTokenAndLoginMethod = async () => { return { token: apiKey, loginMethod: LoginMethod.ANTHROPIC_KEY }; } + if (credentials.loginMethod === LoginMethod.AWS_BEDROCK) { + const accessKeyId = (credentials.secrets as { accessKeyId?: string })?.accessKeyId; + if (!accessKeyId) { + throw new Error('No authentication credentials found'); + } + return { token: accessKeyId, loginMethod: LoginMethod.AWS_BEDROCK }; + } + throw new Error('No authentication credentials found'); }; @@ -480,6 +561,7 @@ const aiStateService = interpret(aiMachine.withConfig({ checkWorkspaceAndToken: checkWorkspaceAndToken, openLogin: openLogin, validateApiKey: validateApiKeyService, + validateAwsCredentials: validateAwsCredentialsService, getTokenAndLoginMethod: getTokenAndLoginMethod, }, actions: { diff --git a/workspaces/mi/mi-extension/src/ai-features/auth.ts b/workspaces/mi/mi-extension/src/ai-features/auth.ts index 304fcccef41..25db68d3b83 100644 --- a/workspaces/mi/mi-extension/src/ai-features/auth.ts +++ b/workspaces/mi/mi-extension/src/ai-features/auth.ts @@ -28,10 +28,11 @@ */ import axios from 'axios'; -import { AIUserToken, AuthCredentials, LoginMethod } from '@wso2/mi-core'; +import { AIUserToken, AuthCredentials, LoginMethod, AwsBedrockSecrets } from '@wso2/mi-core'; import { extension } from '../MIExtensionContext'; import * as vscode from 'vscode'; import { createAnthropic } from '@ai-sdk/anthropic'; +import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock'; import { generateText } from 'ai'; import { CommandIds as PlatformExtCommandIds, IWso2PlatformExtensionAPI } from '@wso2/wso2-platform-core'; import { logInfo, logWarn, logError } from './copilot/logger'; @@ -283,6 +284,9 @@ export const getAccessToken = async (): Promise => { } case LoginMethod.ANTHROPIC_KEY: return credentials.secrets.apiKey; + case LoginMethod.AWS_BEDROCK: + // AWS Bedrock credentials are passed directly to the SDK, not as a single token + return credentials.secrets.accessKeyId; } return undefined; @@ -475,6 +479,97 @@ export const validateApiKey = async (apiKey: string, loginMethod: LoginMethod): } }; +/** + * Validate AWS Bedrock credentials by making a minimal test API call. + */ +export const validateAwsCredentials = async (credentials: { + accessKeyId: string; + secretAccessKey: string; + region: string; + sessionToken?: string; +}): Promise => { + 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 { + logInfo('Validating AWS Bedrock credentials...'); + + const bedrock = createAmazonBedrock({ + region: region, + accessKeyId: accessKeyId, + secretAccessKey: secretAccessKey, + sessionToken: sessionToken, + }); + + // Get regional prefix based on AWS region and construct model ID + const { getBedrockRegionalPrefix } = await import('./connection'); + 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' }] + }); + + logInfo('AWS Bedrock credentials validated successfully'); + + // Store credentials + const authCredentials: AuthCredentials = { + loginMethod: LoginMethod.AWS_BEDROCK, + secrets: { + accessKeyId, + secretAccessKey, + region, + sessionToken + } + }; + await storeAuthCredentials(authCredentials); + + return { token: accessKeyId }; + + } catch (error) { + logError('AWS Bedrock credential validation failed', error); + throw new Error('Validation failed. Please check your AWS credentials and ensure you have access to Amazon Bedrock.'); + } +}; + +/** + * Get stored AWS Bedrock credentials. + */ +export const getAwsBedrockCredentials = async (): Promise => { + const credentials = await getAuthCredentials(); + if (!credentials || credentials.loginMethod !== LoginMethod.AWS_BEDROCK) { + return undefined; + } + return credentials.secrets; +}; + /** * Logout and clear authentication credentials. * Devant session is managed by platform extension separately. diff --git a/workspaces/mi/mi-extension/src/ai-features/connection.ts b/workspaces/mi/mi-extension/src/ai-features/connection.ts index 8d50a89f27d..6b3d138ce66 100644 --- a/workspaces/mi/mi-extension/src/ai-features/connection.ts +++ b/workspaces/mi/mi-extension/src/ai-features/connection.ts @@ -15,8 +15,9 @@ // under the License. import { createAnthropic } from "@ai-sdk/anthropic"; +import { createAmazonBedrock } from "@ai-sdk/amazon-bedrock"; import * as vscode from "vscode"; -import { getAccessToken, getCopilotLlmApiBaseUrl, getLoginMethod, getRefreshedAccessToken } from "./auth"; +import { getAccessToken, getCopilotLlmApiBaseUrl, getLoginMethod, getRefreshedAccessToken, getAwsBedrockCredentials } from "./auth"; import { StateMachineAI, openAIWebview } from "./aiMachine"; import { AI_EVENT_TYPE, LoginMethod } from "@wso2/mi-core"; import { logInfo, logDebug, logError } from "./copilot/logger"; @@ -30,7 +31,29 @@ export type AnthropicModel = | typeof ANTHROPIC_HAIKU_4_5 | typeof ANTHROPIC_SONNET_4_6; +// Bedrock model ID mappings +const BEDROCK_MODEL_MAP: Record = { + [ANTHROPIC_HAIKU_4_5]: "anthropic.claude-3-5-haiku-20241022-v1:0", + [ANTHROPIC_SONNET_4_6]: "anthropic.claude-sonnet-4-20250514-v1:0", +}; + +/** + * Get the regional prefix for Bedrock model IDs based on AWS region. + * Cross-region inference requires a regional prefix (e.g., us., eu.). + */ +export const getBedrockRegionalPrefix = (region: string): string => { + if (region.startsWith('us-')) { + return 'us'; + } else if (region.startsWith('eu-')) { + return 'eu'; + } else if (region.startsWith('ap-')) { + return 'ap'; + } + return 'us'; // default +}; + let cachedAnthropic: ReturnType | null = null; +let cachedBedrock: ReturnType | null = null; let cachedAuthMethod: LoginMethod | null = null; let reLoginPromptInFlight = false; @@ -87,6 +110,7 @@ export async function fetchWithAuth(input: string | URL | Request, options: Requ if (accessToken && loginMethod === LoginMethod.MI_INTEL) { headers['Authorization'] = `Bearer ${accessToken}`; } + // AWS_BEDROCK auth is handled by the AWS SDK directly, no auth headers needed here // Ensure headers object exists // Note: anthropic-beta header for prompt caching is added by the AI SDK @@ -187,6 +211,9 @@ export async function fetchWithAuth(input: string | URL | Request, options: Requ } else if (loginMethod === LoginMethod.ANTHROPIC_KEY) { StateMachineAI.sendEvent(AI_EVENT_TYPE.SILENT_LOGOUT); throw new Error('Authentication failed: Anthropic API key is invalid or expired'); + } else if (loginMethod === LoginMethod.AWS_BEDROCK) { + StateMachineAI.sendEvent(AI_EVENT_TYPE.SILENT_LOGOUT); + throw new Error('Authentication failed: AWS Bedrock credentials are invalid or expired'); } } @@ -226,6 +253,13 @@ export const getAnthropicProvider = async (): Promise> => { + const credentials = await getAwsBedrockCredentials(); + if (!credentials) { + throw new Error("Authentication failed: Unable to get AWS Bedrock credentials"); + } + + // Always recreate to ensure fresh credentials + cachedBedrock = createAmazonBedrock({ + region: credentials.region, + accessKeyId: credentials.accessKeyId, + secretAccessKey: credentials.secretAccessKey, + sessionToken: credentials.sessionToken, + }); + + return cachedBedrock; +}; + export const getAnthropicClient = async (model: AnthropicModel): Promise => { + const loginMethod = await getLoginMethod(); + + if (loginMethod === LoginMethod.AWS_BEDROCK) { + const bedrockProvider = await getBedrockProvider(); + const credentials = await getAwsBedrockCredentials(); + if (!credentials) { + throw new Error("Authentication failed: Unable to get AWS Bedrock credentials"); + } + const regionalPrefix = getBedrockRegionalPrefix(credentials.region); + const bedrockModelId = BEDROCK_MODEL_MAP[model]; + if (!bedrockModelId) { + throw new Error(`No Bedrock model mapping found for: ${model}`); + } + const fullModelId = `${regionalPrefix}.${bedrockModelId}`; + logDebug(`Using Bedrock model: ${fullModelId}`); + return bedrockProvider(fullModelId); + } + const provider = await getAnthropicProvider(); return provider(model); }; /** * Returns cache control options for prompt caching - * @returns Cache control options for Anthropic + * @returns Cache control options for Anthropic or Bedrock */ -export const getProviderCacheControl = () => { +export const getProviderCacheControl = async (): Promise> => { + const loginMethod = await getLoginMethod(); + if (loginMethod === LoginMethod.AWS_BEDROCK) { + return { bedrock: { cacheControl: { type: "ephemeral" } } }; + } return { anthropic: { cacheControl: { type: "ephemeral" } } }; }; diff --git a/workspaces/mi/mi-extension/src/rpc-managers/ai-features/rpc-manager.ts b/workspaces/mi/mi-extension/src/rpc-managers/ai-features/rpc-manager.ts index 61855238950..cc0ceb38428 100644 --- a/workspaces/mi/mi-extension/src/rpc-managers/ai-features/rpc-manager.ts +++ b/workspaces/mi/mi-extension/src/rpc-managers/ai-features/rpc-manager.ts @@ -478,11 +478,11 @@ export class MIAIPanelRpcManager implements MIAIPanelAPI { } /** - * Check if user is using their own Anthropic API key + * Check if user is using their own Anthropic API key or AWS Bedrock credentials */ async hasAnthropicApiKey(): Promise { const loginMethod = await getLoginMethod(); - return loginMethod === LoginMethod.ANTHROPIC_KEY; + return loginMethod === LoginMethod.ANTHROPIC_KEY || loginMethod === LoginMethod.AWS_BEDROCK; } /** diff --git a/workspaces/mi/mi-visualizer/src/views/AIPanel/component/AIChatHeader.tsx b/workspaces/mi/mi-visualizer/src/views/AIPanel/component/AIChatHeader.tsx index f4e2f962aa8..cc283f9efdb 100644 --- a/workspaces/mi/mi-visualizer/src/views/AIPanel/component/AIChatHeader.tsx +++ b/workspaces/mi/mi-visualizer/src/views/AIPanel/component/AIChatHeader.tsx @@ -18,6 +18,7 @@ import React, { useState, useEffect } from "react"; import { Button, Codicon } from "@wso2/ui-toolkit"; +import { LoginMethod } from "@wso2/mi-core"; import { Badge, Header, HeaderButtons, ResetsInBadge } from '../styles'; import { useMICopilotContext } from "./MICopilotContext"; import SessionSwitcher from "./SessionSwitcher"; @@ -42,6 +43,7 @@ const AIChatHeader: React.FC = () => { deleteSession } = useMICopilotContext(); const [hasApiKey, setHasApiKey] = useState(false); + const [isAwsBedrock, setIsAwsBedrock] = useState(false); const handleLogout = async () => { await rpcClient?.getMiDiagramRpcClient().logoutFromMIAccount(); @@ -50,6 +52,9 @@ const AIChatHeader: React.FC = () => { const checkApiKey = async () => { const hasApiKey = await rpcClient?.getMiAiPanelRpcClient().hasAnthropicApiKey(); setHasApiKey(hasApiKey); + // Check if specifically using AWS Bedrock + const machineView = await rpcClient?.getAIVisualizerState(); + setIsAwsBedrock(machineView?.loginMethod === LoginMethod.AWS_BEDROCK); }; // Check for API key on component mount @@ -66,9 +71,9 @@ const AIChatHeader: React.FC = () => {
- Copilot is using your API Key + {isAwsBedrock ? "Copilot is using your AWS Bedrock Account" : "Copilot is using your API Key"}
- Logout to clear the API key + {isAwsBedrock ? "Logout to clear the credentials" : "Logout to clear the API key"}
) : ( <> diff --git a/workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx b/workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx index be6a870b4ae..8ad8517e544 100644 --- a/workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx +++ b/workspaces/mi/mi-visualizer/src/views/AIPanel/component/WaitingForLoginSection.tsx @@ -194,6 +194,15 @@ export const WaitingForLoginSection = ({ loginMethod, isValidating = false, erro const [apiKey, setApiKey] = useState(""); const [showApiKey, setShowApiKey] = useState(false); const [clientError, setClientError] = useState(""); + const [awsCredentials, setAwsCredentials] = useState({ + accessKeyId: "", + secretAccessKey: "", + region: "", + sessionToken: "" + }); + const [showAwsSecretKey, setShowAwsSecretKey] = useState(false); + const [showAwsAccessKey, setShowAwsAccessKey] = useState(false); + const [showAwsSessionToken, setShowAwsSessionToken] = useState(false); const cancelLogin = () => { rpcClient.sendAIStateEvent(AI_EVENT_TYPE.CANCEL_LOGIN); @@ -241,6 +250,42 @@ export const WaitingForLoginSection = ({ loginMethod, isValidating = false, erro const toggleApiKeyVisibility = () => { setShowApiKey(!showApiKey); }; + + const handleAwsCredentialChange = (field: string) => (e: any) => { + const value = e.target?.value ?? ''; + setAwsCredentials(prev => ({ ...prev, [field]: value })); + if (clientError) { + setClientError(""); + } + }; + + const connectWithAwsBedrock = () => { + setClientError(""); + const { accessKeyId, secretAccessKey, region, sessionToken } = awsCredentials; + + if (!accessKeyId.trim()) { + setClientError("Please enter your AWS Access Key ID"); + return; + } + if (!secretAccessKey.trim()) { + setClientError("Please enter your AWS Secret Access Key"); + return; + } + if (!region.trim()) { + setClientError("Please enter your AWS Region"); + return; + } + + rpcClient.sendAIStateEvent({ + type: AI_EVENT_TYPE.SUBMIT_AWS_CREDENTIALS, + payload: { + accessKeyId: accessKeyId.trim(), + secretAccessKey: secretAccessKey.trim(), + region: region.trim(), + sessionToken: sessionToken.trim() || undefined + }, + } as any); + }; // Show either server error (from validation) or client error (from form validation) const displayError = errorMessage || clientError; @@ -303,6 +348,120 @@ export const WaitingForLoginSection = ({ loginMethod, isValidating = false, erro ); } + if (loginMethod === LoginMethod.AWS_BEDROCK) { + return ( + + + Connect with AWS Bedrock + + Enter your AWS credentials to connect to MI Copilot via Amazon Bedrock. + Your credentials will be securely stored and used for authentication. + + + {/* AWS Access Key ID */} + + + + setShowAwsAccessKey(!showAwsAccessKey)} + title={showAwsAccessKey ? "Hide" : "Show"} + {...(isValidating ? { disabled: true } : {})} + > + + + + + + {/* AWS Secret Access Key */} + + + + setShowAwsSecretKey(!showAwsSecretKey)} + title={showAwsSecretKey ? "Hide" : "Show"} + {...(isValidating ? { disabled: true } : {})} + > + + + + + + {/* AWS Region */} + + + + + + + {/* Session Token (Optional) */} + + + + setShowAwsSessionToken(!showAwsSessionToken)} + title={showAwsSessionToken ? "Hide" : "Show"} + {...(isValidating ? { disabled: true } : {})} + > + + + + + + {displayError && ( + + + {displayError} + + )} + + + + {isValidating ? "Validating..." : "Connect with AWS Bedrock"} + + + Cancel + + + + + ); + } + // Default: MI_INTEL login method return ( diff --git a/workspaces/mi/mi-visualizer/src/views/AIPanel/index.tsx b/workspaces/mi/mi-visualizer/src/views/AIPanel/index.tsx index 6d6e55ca812..f2fcaf83eeb 100644 --- a/workspaces/mi/mi-visualizer/src/views/AIPanel/index.tsx +++ b/workspaces/mi/mi-visualizer/src/views/AIPanel/index.tsx @@ -57,7 +57,7 @@ export const AIPanel = () => { const authenticatingState = (state as any).Authenticating; // Determine if we're validating - const isValidating = authenticatingState === 'validatingApiKey'; + const isValidating = authenticatingState === 'validatingApiKey' || authenticatingState === 'validatingAwsCredentials'; // Show the appropriate form based on the substate if (authenticatingState === 'apiKeyFlow' || authenticatingState === 'validatingApiKey') { @@ -68,6 +68,14 @@ export const AIPanel = () => { errorMessage={errorMessage} /> ); + } else if (authenticatingState === 'awsBedrockFlow' || authenticatingState === 'validatingAwsCredentials') { + setViewComponent( + + ); } else { // For ssoFlow or determineFlow, show waiting message setViewComponent( diff --git a/workspaces/mi/mi-visualizer/src/views/LoggedOutWindow/index.tsx b/workspaces/mi/mi-visualizer/src/views/LoggedOutWindow/index.tsx index 284ad692807..a6b3d28d619 100644 --- a/workspaces/mi/mi-visualizer/src/views/LoggedOutWindow/index.tsx +++ b/workspaces/mi/mi-visualizer/src/views/LoggedOutWindow/index.tsx @@ -140,6 +140,10 @@ export const SignInToCopilotMessage = (props: { showProjectHeader?: boolean }) = rpcClient.sendAIStateEvent(AI_EVENT_TYPE.AUTH_WITH_API_KEY); }; + const handleAwsBedrockClick = () => { + rpcClient.sendAIStateEvent(AI_EVENT_TYPE.AUTH_WITH_AWS_BEDROCK); + }; + return ( @@ -168,6 +172,7 @@ export const SignInToCopilotMessage = (props: { showProjectHeader?: boolean }) = Login to MI Copilot or Enter your Anthropic API key + Enter your AWS Bedrock credentials ); From ab6392e68af2401102ed11d396ddae96d349da1c Mon Sep 17 00:00:00 2001 From: Ravindu Wegiriya Date: Wed, 25 Feb 2026 13:44:29 +0530 Subject: [PATCH 2/3] Do reviewed changes --- .../src/ai-features/connection.ts | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/workspaces/mi/mi-extension/src/ai-features/connection.ts b/workspaces/mi/mi-extension/src/ai-features/connection.ts index 6b3d138ce66..c4a28ad7729 100644 --- a/workspaces/mi/mi-extension/src/ai-features/connection.ts +++ b/workspaces/mi/mi-extension/src/ai-features/connection.ts @@ -42,14 +42,19 @@ const BEDROCK_MODEL_MAP: Record = { * Cross-region inference requires a regional prefix (e.g., us., eu.). */ export const getBedrockRegionalPrefix = (region: string): string => { - if (region.startsWith('us-')) { - return 'us'; - } else if (region.startsWith('eu-')) { - return 'eu'; - } else if (region.startsWith('ap-')) { - return 'ap'; + const prefix = region.split('-')[0]; + switch (prefix) { + case 'us': + case 'eu': + case 'ap': + case 'ca': + case 'sa': + case 'me': + case 'af': + return prefix; + default: + return 'us'; } - return 'us'; // default }; let cachedAnthropic: ReturnType | null = null; From 05f707aa68587a93a81e9add0c704f84752b5854 Mon Sep 17 00:00:00 2001 From: Ravindu Wegiriya Date: Wed, 25 Feb 2026 13:44:44 +0530 Subject: [PATCH 3/3] do reviewed changes --- .../mi/mi-extension/src/ai-features/aiMachine.ts | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts b/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts index cba5bb1c676..b6f44bd53d6 100644 --- a/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts +++ b/workspaces/mi/mi-extension/src/ai-features/aiMachine.ts @@ -470,9 +470,9 @@ const checkWorkspaceAndToken = async (): Promise<{ workspaceSupported: boolean; tokenData = { token: apiKey, loginMethod: LoginMethod.ANTHROPIC_KEY }; } } else if (credentials?.loginMethod === LoginMethod.AWS_BEDROCK) { - const accessKeyId = (credentials.secrets as { accessKeyId?: string })?.accessKeyId; - if (accessKeyId) { - tokenData = { token: accessKeyId, loginMethod: LoginMethod.AWS_BEDROCK }; + const secrets = credentials.secrets as { accessKeyId?: string; secretAccessKey?: string; region?: string }; + if (secrets.accessKeyId && secrets.secretAccessKey && secrets.region) { + tokenData = { token: secrets.accessKeyId, loginMethod: LoginMethod.AWS_BEDROCK }; } } @@ -546,11 +546,11 @@ const getTokenAndLoginMethod = async () => { } if (credentials.loginMethod === LoginMethod.AWS_BEDROCK) { - const accessKeyId = (credentials.secrets as { accessKeyId?: string })?.accessKeyId; - if (!accessKeyId) { - throw new Error('No authentication credentials found'); + const secrets = credentials.secrets as { accessKeyId?: string; secretAccessKey?: string; region?: string }; + if (!secrets.accessKeyId || !secrets.secretAccessKey || !secrets.region) { + throw new Error('Incomplete AWS Bedrock credentials. Please log in again.'); } - return { token: accessKeyId, loginMethod: LoginMethod.AWS_BEDROCK }; + return { token: secrets.accessKeyId, loginMethod: LoginMethod.AWS_BEDROCK }; } throw new Error('No authentication credentials found');