Skip to content

Add AWS Bedrock Login #1555

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 3 commits into
base: main
Choose a base branch
from
Open

Add AWS Bedrock Login #1555

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7fd09f3
Add AWS Bedrock Login
ravindu25 Feb 25, 2026
ab6392e
Do reviewed changes
ravindu25 Feb 25, 2026
05f707a
do reviewed changes
ravindu25 Feb 25, 2026
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
30 changes: 28 additions & 2 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 @@ -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)
Expand All @@ -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",
Expand All @@ -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;
Expand All @@ -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 {
Expand All @@ -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;
Expand All @@ -194,6 +211,10 @@ export type AuthCredentials =
| {
loginMethod: LoginMethod.ANTHROPIC_KEY;
secrets: AnthropicKeySecrets;
}
| {
loginMethod: LoginMethod.AWS_BEDROCK;
secrets: AwsBedrockSecrets;
};

export interface AIUserToken {
Expand Down
1 change: 1 addition & 0 deletions workspaces/mi/mi-extension/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
82 changes: 82 additions & 0 deletions workspaces/mi/mi-extension/src/ai-features/aiMachine.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import {
initiateInbuiltAuth,
logout,
validateApiKey,
validateAwsCredentials,
isPlatformExtensionAvailable,
isDevantUserLoggedIn,
getPlatformStsToken,
Expand Down Expand Up @@ -150,6 +151,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 @@ -166,6 +173,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 @@ -253,6 +264,48 @@ 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,
})
},
[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'
})
}
}
}
}
},
Expand Down Expand Up @@ -324,6 +377,13 @@ const aiMachine = createMachine<AIMachineContext, AIMachineSendableEvent>({
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({
Expand Down Expand Up @@ -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 };
Expand Down Expand Up @@ -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) {
Expand All @@ -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');
};

Expand All @@ -480,6 +561,7 @@ const aiStateService = interpret(aiMachine.withConfig({
checkWorkspaceAndToken: checkWorkspaceAndToken,
openLogin: openLogin,
validateApiKey: validateApiKeyService,
validateAwsCredentials: validateAwsCredentialsService,
getTokenAndLoginMethod: getTokenAndLoginMethod,
},
actions: {
Expand Down
Loading
Loading