Skip to content

Commit 57e69af

Browse files
samclark2015samclark2015
committed
Better flow for saving from config modal
1 parent 61a7d1c commit 57e69af

File tree

3 files changed

+104
-73
lines changed

3 files changed

+104
-73
lines changed

‎extensions/positron-assistant/src/config.ts‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -252,16 +252,17 @@ async function saveModel(userConfig: positron.ai.LanguageModelConfig, sources: p
252252
...otherConfig,
253253
};
254254

255-
// Update global state
256-
await context.globalState.update(
257-
'positron.assistant.models',
258-
[...existingConfigs, newConfig]
259-
);
260255

261256
// Register the new model FIRST, before saving configuration
262257
try {
263258
await registerModel(newConfig, context, storage);
264259

260+
// Update global state
261+
await context.globalState.update(
262+
'positron.assistant.models',
263+
[...existingConfigs, newConfig]
264+
);
265+
265266
positron.ai.addLanguageModelConfig(expandConfigToSource(newConfig));
266267

267268
// Refresh CopilotService signed-in state if this is a copilot model

‎extensions/positron-assistant/src/extension.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,12 +80,15 @@ export const log = vscode.window.createOutputChannel('Assistant', { log: true })
8080

8181
export async function registerModel(config: StoredModelConfig, context: vscode.ExtensionContext, storage: SecretStorage) {
8282
try {
83-
const modelConfig = await getModelConfiguration(config.id, context, storage);
83+
const modelConfig: ModelConfig = {
84+
...config,
85+
apiKey: undefined // will be filled in below if needed
86+
};
8487

8588
if (modelConfig?.baseUrl) {
8689
const apiKey = await storage.get(`apiKey-${modelConfig.id}`);
8790
if (apiKey) {
88-
(modelConfig as any).apiKey = apiKey;
91+
modelConfig.apiKey = apiKey;
8992
}
9093
}
9194

‎extensions/positron-assistant/src/models.ts‎

Lines changed: 93 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import * as vscode from 'vscode';
77
import * as positron from 'positron';
88
import * as ai from 'ai';
9-
import { expandConfigToSource, getMaxConnectionAttempts, getProviderTimeoutMs, ModelConfig, SecretStorage } from './config';
9+
import { expandConfigToSource, getMaxConnectionAttempts, getProviderTimeoutMs, getStoredModels, ModelConfig, SecretStorage } from './config';
1010
import { AnthropicProvider, createAnthropic } from '@ai-sdk/anthropic';
1111
import { AzureOpenAIProvider, createAzure } from '@ai-sdk/azure';
1212
import { createVertex, GoogleVertexProvider } from '@ai-sdk/google-vertex';
@@ -377,7 +377,7 @@ abstract class AILanguageModel implements positron.ai.LanguageModelChatProvider
377377
} catch (error) {
378378
const messagePrefix = `[${this.providerName}] '${model}'`;
379379
log.warn(`${messagePrefix} Error sending test message: ${JSON.stringify(error, null, 2)}`);
380-
const errorMsg = this.parseProviderError(error) ||
380+
const errorMsg = await this.parseProviderError(error) ||
381381
(ai.AISDKError.isInstance(error) ? error.message : JSON.stringify(error, null, 2));
382382
errors.push(errorMsg);
383383
}
@@ -524,7 +524,7 @@ abstract class AILanguageModel implements positron.ai.LanguageModelChatProvider
524524
flushAccumulatedTextDeltas();
525525
const messagePrefix = `[${this.providerName}] [${model.name}]'`;
526526
log.warn(`${messagePrefix} RECV error: ${JSON.stringify(part.error, null, 2)}`);
527-
const errorMsg = this.parseProviderError(part.error) ||
527+
const errorMsg = await this.parseProviderError(part.error) ||
528528
(typeof part.error === 'string' ? part.error : JSON.stringify(part.error, null, 2));
529529
return new Error(`${messagePrefix} Error sending test message: ${errorMsg}`);
530530
}
@@ -592,7 +592,7 @@ abstract class AILanguageModel implements positron.ai.LanguageModelChatProvider
592592
* @param error The error object returned by the provider.
593593
* @returns A user-friendly error message or undefined if not specifically handled.
594594
*/
595-
parseProviderError(error: any): string | undefined {
595+
async parseProviderError(error: any): Promise<string | undefined> {
596596
// Try to extract an API error message with ai-sdk
597597
if (ai.APICallError.isInstance(error)) {
598598
const responseBody = error.responseBody;
@@ -1046,6 +1046,10 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
10461046
bedrockClient: BedrockClient;
10471047
inferenceProfiles: InferenceProfileSummary[] = [];
10481048

1049+
// Keep state while we're resolving the connection
1050+
// Used to adjust error handling for SSO login prompts
1051+
private _resolvingConnection: boolean = false;
1052+
10491053
constructor(_config: ModelConfig, _context?: vscode.ExtensionContext, _storage?: SecretStorage) {
10501054
// Update a stale model configuration to the latest defaults
10511055
const models = availableModels.get('amazon-bedrock')?.map(m => m.identifier) || [];
@@ -1090,8 +1094,8 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
10901094
* @param error The error object
10911095
* @returns A user-friendly error message or undefined if not specifically handled.
10921096
*/
1093-
override parseProviderError(error: any): string | undefined {
1094-
const aiSdkError = super.parseProviderError(error);
1097+
override async parseProviderError(error: any): Promise<string | undefined> {
1098+
const aiSdkError = await super.parseProviderError(error);
10951099
if (aiSdkError) {
10961100
return aiSdkError;
10971101
}
@@ -1104,75 +1108,44 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
11041108
const message = error.message;
11051109

11061110
if (!message) {
1107-
return super.parseProviderError(error);
1111+
return await super.parseProviderError(error);
11081112
}
11091113

11101114
if (name === 'CredentialsProviderError') {
11111115
// This is specifically the error thrown the refresh token is expired and
11121116
// we need to re-authenticate with AWS SSO
11131117
if (message.includes('aws sso login')) {
11141118
// Give the user the option to login automatically
1115-
// Display an error message with an action the user can take
1116-
const isConnectionTest = error.stack?.includes('resolveConnection');
1117-
const action = { title: vscode.l10n.t('Run in Terminal'), id: 'aws-sso-login' };
1118-
vscode.window.showErrorMessage(`Amazon Bedrock: ${message}`, action).then(async selection => {
1119-
if (selection?.id === action.id) {
1120-
// User chose to login, so we need to refresh the credentials
1121-
1122-
// Grab the profile & region to refresh from the Bedrock client config
1123-
const profile = this.bedrockClient.config.profile;
1124-
// Region may be an async function or a string, so handle both cases
1125-
const region = typeof this.bedrockClient.config.region === 'function' ? await this.bedrockClient.config.region() : this.bedrockClient.config.region;
1126-
// Execute the AWS SSO login command as a native task
1127-
const taskExecution = await vscode.tasks.executeTask(new vscode.Task(
1128-
{ type: 'shell' },
1129-
vscode.TaskScope.Workspace,
1130-
'AWS SSO Login',
1131-
'AWS',
1132-
new vscode.ShellExecution(`aws sso login --profile ${profile} --region ${region}`)
1133-
));
1134-
1135-
vscode.tasks.onDidEndTaskProcess(e => {
1136-
if (e.execution === taskExecution) {
1137-
// Notify the user of the result
1138-
const success = e.exitCode === 0 || e.exitCode === undefined;
1139-
if (success) {
1140-
// Success
1141-
vscode.window.showInformationMessage(vscode.l10n.t('AWS login completed successfully'));
1142-
} else {
1143-
// Failure
1144-
vscode.window.showErrorMessage(vscode.l10n.t('AWS login failed with exit code {0}', e.exitCode));
1145-
}
1146-
1147-
// Open a URI to bring Positron to the foreground
1148-
// This is a little sneaky, but works + no other native method
1149-
const redirectUri = vscode.Uri.from({ scheme: vscode.env.uriScheme });
1150-
vscode.env.openExternal(redirectUri);
1151-
1152-
if (success && isConnectionTest) {
1153-
// If we were in a connection test, re-run it now that we've logged in
1154-
registerModelWithAPI(
1155-
this._config,
1156-
this._context,
1157-
this._storage,
1158-
this
1159-
).then(() => {
1160-
positron.ai.addLanguageModelConfig(expandConfigToSource(this._config));
1161-
PositronAssistantApi.get().notifySignIn(this._config.name);
1162-
});
1163-
}
1164-
}
1165-
});
1119+
const existingModels = getStoredModels(this._context);
1120+
1121+
// Check if our model is already registered
1122+
if (!existingModels.some(m => m.provider === this._config.provider)) {
1123+
// The model is not yet registered, so just refresh without prompting
1124+
if (await this.refreshCredentials(true)) {
1125+
// If we're successful, return undefined to indicate no error
1126+
return undefined;
11661127
}
1167-
});
11681128

1169-
if (isConnectionTest) {
1170-
// We're in a connection test, so throw an AssistantError to avoid showing a message box
1171-
// but that stops the model provider from being registered in core
1172-
throw new AssistantError(message, false);
11731129
} else {
1174-
// We are in a chat response, so we should return an error to display in the chat pane
1175-
throw new Error(vscode.l10n.t(`AWS login required. Please run \`aws sso login --profile ${this.bedrockClient.config.profile} --region ${this.bedrockClient.config.region}\` in the terminal, and retry this request.`));
1130+
// The model has already been registered, so we can prompt the user to login
1131+
// Display an error message with an action the user can take
1132+
const isConnectionTest = this._resolvingConnection;
1133+
const action = { title: vscode.l10n.t('Run in Terminal'), id: 'aws-sso-login' };
1134+
vscode.window.showErrorMessage(`Amazon Bedrock: ${message}`, action).then(async selection => {
1135+
if (selection?.id === action.id) {
1136+
// User chose to login, so we need to refresh the credentials
1137+
await this.refreshCredentials(isConnectionTest);
1138+
}
1139+
});
1140+
1141+
if (isConnectionTest) {
1142+
// We're in a connection test, so throw an AssistantError to avoid showing a message box
1143+
// but that stops the model provider from being registered in core
1144+
throw new AssistantError(message, false);
1145+
} else {
1146+
// We are in a chat response, so we should return an error to display in the chat pane
1147+
throw new Error(vscode.l10n.t(`AWS login required. Please run \`aws sso login --profile ${this.bedrockClient.config.profile} --region ${this.bedrockClient.config.region}\` in the terminal, and retry this request.`));
1148+
}
11761149
}
11771150
} else {
11781151
return vscode.l10n.t(`Invalid AWS credentials. {0}`, message);
@@ -1182,19 +1155,73 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
11821155
return vscode.l10n.t(`Amazon Bedrock error: {0}`, message);
11831156
}
11841157

1158+
private async refreshCredentials(reregister: boolean = false): Promise<boolean> {
1159+
// Grab the profile & region to refresh from the Bedrock client config
1160+
const profile = this.bedrockClient.config.profile;
1161+
// Region may be an async function or a string, so handle both cases
1162+
const region = typeof this.bedrockClient.config.region === 'function' ? await this.bedrockClient.config.region() : this.bedrockClient.config.region;
1163+
// Execute the AWS SSO login command as a native task
1164+
const taskExecution = await vscode.tasks.executeTask(new vscode.Task(
1165+
{ type: 'shell' },
1166+
vscode.TaskScope.Workspace,
1167+
'AWS SSO Login',
1168+
'AWS',
1169+
new vscode.ShellExecution(`aws sso login --profile ${profile} --region ${region}`)
1170+
));
1171+
1172+
const result = new Promise<boolean>((resolve) => {
1173+
vscode.tasks.onDidEndTaskProcess(e => {
1174+
if (e.execution === taskExecution) {
1175+
// Notify the user of the result
1176+
const success = e.exitCode === 0 || e.exitCode === undefined;
1177+
if (success) {
1178+
// Success
1179+
vscode.window.showInformationMessage(vscode.l10n.t('AWS login completed successfully'));
1180+
} else {
1181+
// Failure
1182+
vscode.window.showErrorMessage(vscode.l10n.t('AWS login failed with exit code {0}', e.exitCode));
1183+
}
1184+
1185+
// Open a URI to bring Positron to the foreground
1186+
// This is a little sneaky, but works + no other native method
1187+
const redirectUri = vscode.Uri.from({ scheme: vscode.env.uriScheme });
1188+
vscode.env.openExternal(redirectUri);
1189+
1190+
if (success && reregister) {
1191+
// If we were in a connection test, re-run it now that we've logged in
1192+
registerModelWithAPI(
1193+
this._config,
1194+
this._context,
1195+
this._storage,
1196+
this
1197+
).then(() => {
1198+
positron.ai.addLanguageModelConfig(expandConfigToSource(this._config));
1199+
PositronAssistantApi.get().notifySignIn(this._config.name);
1200+
});
1201+
}
1202+
resolve(success);
1203+
}
1204+
});
1205+
});
1206+
return result;
1207+
}
1208+
11851209
override async resolveConnection(token: vscode.CancellationToken): Promise<Error | undefined> {
11861210
// The Vercel and Bedrock SDKs both use the node provider chain for credentials so getting a listing
11871211
// means the credentials are valid.
11881212
log.debug(`[${this.providerName}] Resolving connection by fetching available models...`);
1213+
this._resolvingConnection = true;
11891214
try {
11901215
await this.resolveModels(token);
11911216
} catch (error) {
11921217
// Try to parse specific Bedrock errors
11931218
// This way, we can handle SSO login errors specifically
1194-
const parsedError = this.parseProviderError(error);
1219+
const parsedError = await this.parseProviderError(error);
11951220
if (parsedError) {
11961221
return new Error(parsedError);
11971222
}
1223+
} finally {
1224+
this._resolvingConnection = false;
11981225
}
11991226

12001227
return undefined;

0 commit comments

Comments
 (0)