#27によって発生していた画面リロード時に認証情報が表示されない問題、LLMの一覧取得ができない問題など種々の課題を修正

Author: daisuke-awaji

src/main/api/bedrock/services/imageService.ts

@@ -1,5 +1,5 @@
 import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime'
-import type { AWSCredentials, ServiceContext } from '../types'
+import type { ServiceContext } from '../types'
 import type {
   AspectRatio,
   GenerateImageRequest,
@@ -111,14 +111,20 @@ export class ImageService {
   ]
 
   constructor(private context: ServiceContext) {
-    const awsCredentials: AWSCredentials = this.context.store.get('aws')
-    const { credentials, region } = awsCredentials
+    const { accessKeyId, secretAccessKey, sessionToken, region } = this.context.store.get('aws')
 
-    if (!credentials || !credentials.accessKeyId || !credentials.secretAccessKey || !region) {
+    if (!accessKeyId || !secretAccessKey || !region) {
       console.warn('AWS credentials not configured')
     }
 
-    this.runtimeClient = new BedrockRuntimeClient(awsCredentials)
+    this.runtimeClient = new BedrockRuntimeClient({
+      credentials: {
+        accessKeyId,
+        secretAccessKey,
+        sessionToken
+      },
+      region
+    })
   }
 
   private getModelType(modelId: ImageGeneratorModel): ModelType {

src/main/api/bedrock/services/modelService.ts

@@ -1,6 +1,6 @@
 import { getDefaultPromptRouter, getModelsForRegion } from '../models'
 import { getAccountId } from '../utils/awsUtils'
-import type { AWSCredentials, ServiceContext } from '../types'
+import type { ServiceContext } from '../types'
 import { BedrockSupportRegion } from '../../../../types/llm'
 
 export class ModelService {
@@ -10,14 +10,14 @@ export class ModelService {
   constructor(private context: ServiceContext) {}
 
   async listModels() {
-    const awsCredentials: AWSCredentials = this.context.store.get('aws')
-    const { credentials, region } = awsCredentials
-    if (!credentials || !credentials.accessKeyId || !credentials.secretAccessKey || !region) {
+    const { accessKeyId, secretAccessKey, sessionToken, region } = this.context.store.get('aws')
+
+    if (!accessKeyId || !secretAccessKey || !region) {
       console.warn('AWS credentials not configured')
       return []
     }
 
-    const cacheKey = `${region}-${credentials.accessKeyId}`
+    const cacheKey = `${region}-${accessKeyId}`
     const cachedData = this.modelCache[cacheKey]
 
     if (
@@ -31,7 +31,12 @@ export class ModelService {
     try {
       const models = getModelsForRegion(region as BedrockSupportRegion)
 
-      const accountId = await getAccountId(awsCredentials)
+      const accountId = await getAccountId({
+        accessKeyId,
+        secretAccessKey,
+        sessionToken,
+        region
+      })
       const promptRouterModels = accountId ? getDefaultPromptRouter(accountId, region) : []
       const result = [...models, ...promptRouterModels]
       this.modelCache[cacheKey] = [...result, { _timestamp: Date.now() } as any]

src/main/api/bedrock/types.ts

@@ -1,5 +1,4 @@
 import { Message } from '@aws-sdk/client-bedrock-runtime'
-import type { AwsCredentialIdentity } from '@smithy/types'
 
 export type CallConverseAPIProps = {
   modelId: string
@@ -9,7 +8,9 @@ export type CallConverseAPIProps = {
 }
 
 export type AWSCredentials = {
-  credentials: AwsCredentialIdentity
+  accessKeyId: string
+  secretAccessKey: string
+  sessionToken?: string
   region: string
 }
 

src/main/api/bedrock/utils/awsUtils.ts

@@ -52,7 +52,14 @@ export function getAlternateRegionOnThrottling(
 
 export async function getAccountId(awsCredentials: AWSCredentials) {
   try {
-    const sts = new STSClient(awsCredentials)
+    const sts = new STSClient({
+      credentials: {
+        accessKeyId: awsCredentials.accessKeyId,
+        secretAccessKey: awsCredentials.secretAccessKey,
+        sessionToken: awsCredentials?.sessionToken
+      },
+      region: awsCredentials.region
+    })
     const command = new GetCallerIdentityCommand({})
     const res = await sts.send(command)
     return res.Account

src/preload/store.ts

@@ -37,6 +37,7 @@ type StoreScheme = {
     region: string
     accessKeyId: string
     secretAccessKey: string
+    sessionToken?: string
   }
   customAgents: CustomAgent[]
   selectedAgentId: string

src/renderer/src/contexts/SettingsContext.tsx

@@ -294,6 +294,7 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
       setStateAwsAccessKeyId(awsConfig.accessKeyId || '')
       setStateAwsSecretAccessKey(awsConfig.secretAccessKey || '')
       setStateAwsSessionToken(awsConfig.sessionToken || '')
+      console.log({ awsConfig })
     }
 
     // Load Custom Agents
@@ -516,7 +517,7 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
 
   const saveAwsConfig = (credentials: AwsCredentialIdentity, region: string) => {
     window.store.set('aws', {
-      accessKey: credentials.accessKeyId,
+      accessKeyId: credentials.accessKeyId,
       secretAccessKey: credentials.secretAccessKey,
       sessionToken: credentials.sessionToken,
       region