Skip to content

Commit 4eba0bc

Browse files
daisuke-awajidaisuke-awaji
authored
Merge pull request #27 from inoue22/main
STSセッショントークンのサポート（型定義の強化）
2 parents 6c3a6c3 + 4b5aa14 commit 4eba0bc

File tree

13 files changed

+14530
-83
lines changed

13 files changed

+14530
-83
lines changed

‎package-lock.json‎

Lines changed: 14441 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
"@electron-toolkit/eslint-config-prettier": "2.0.0",
6868
"@electron-toolkit/eslint-config-ts": "2.0.0",
6969
"@electron-toolkit/tsconfig": "1.0.1",
70+
"@smithy/types": "4.1.0",
7071
"@types/diff": "5.2.3",
7172
"@types/express": "4.17.21",
7273
"@types/jest": "29.5.12",

‎src/main/api/bedrock/client.ts‎

Lines changed: 6 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,35 +3,14 @@ import { BedrockClient } from '@aws-sdk/client-bedrock'
33
import { BedrockAgentRuntimeClient } from '@aws-sdk/client-bedrock-agent-runtime'
44
import type { AWSCredentials } from './types'
55

6-
export function createRuntimeClient(credentials: AWSCredentials) {
7-
const { region, accessKeyId, secretAccessKey } = credentials
8-
return new BedrockRuntimeClient({
9-
credentials: {
10-
accessKeyId,
11-
secretAccessKey
12-
},
13-
region
14-
})
6+
export function createRuntimeClient(awsCredentials: AWSCredentials) {
7+
return new BedrockRuntimeClient(awsCredentials)
158
}
169

17-
export function createBedrockClient(credentials: AWSCredentials) {
18-
const { region, accessKeyId, secretAccessKey } = credentials
19-
return new BedrockClient({
20-
credentials: {
21-
accessKeyId,
22-
secretAccessKey
23-
},
24-
region
25-
})
10+
export function createBedrockClient(awsCredentials: AWSCredentials) {
11+
return new BedrockClient(awsCredentials)
2612
}
2713

28-
export function createAgentRuntimeClient(credentials: AWSCredentials) {
29-
const { region, accessKeyId, secretAccessKey } = credentials
30-
return new BedrockAgentRuntimeClient({
31-
credentials: {
32-
accessKeyId,
33-
secretAccessKey
34-
},
35-
region
36-
})
14+
export function createAgentRuntimeClient(awsCredentials: AWSCredentials) {
15+
return new BedrockAgentRuntimeClient(awsCredentials)
3716
}

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

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime'
2-
import type { ServiceContext } from '../types'
2+
import type { AWSCredentials, ServiceContext } from '../types'
33
import type {
44
AspectRatio,
55
GenerateImageRequest,
@@ -111,18 +111,14 @@ export class ImageService {
111111
]
112112

113113
constructor(private context: ServiceContext) {
114-
const { region, accessKeyId, secretAccessKey } = this.context.store.get('aws')
115-
if (!accessKeyId || !secretAccessKey) {
114+
const awsCredentials: AWSCredentials = this.context.store.get('aws')
115+
const { credentials, region } = awsCredentials
116+
117+
if (!credentials || !credentials.accessKeyId || !credentials.secretAccessKey || !region) {
116118
console.warn('AWS credentials not configured')
117119
}
118120

119-
this.runtimeClient = new BedrockRuntimeClient({
120-
credentials: {
121-
accessKeyId,
122-
secretAccessKey
123-
},
124-
region
125-
})
121+
this.runtimeClient = new BedrockRuntimeClient(awsCredentials)
126122
}
127123

128124
private getModelType(modelId: ImageGeneratorModel): ModelType {

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { getDefaultPromptRouter, getModelsForRegion } from '../models'
22
import { getAccountId } from '../utils/awsUtils'
3-
import type { ServiceContext } from '../types'
3+
import type { AWSCredentials, ServiceContext } from '../types'
44
import { BedrockSupportRegion } from '../../../../types/llm'
55

66
export class ModelService {
@@ -10,14 +10,14 @@ export class ModelService {
1010
constructor(private context: ServiceContext) {}
1111

1212
async listModels() {
13-
const credentials = this.context.store.get('aws')
14-
const { region, accessKeyId, secretAccessKey } = credentials
15-
if (!region || !accessKeyId || !secretAccessKey) {
13+
const awsCredentials: AWSCredentials = this.context.store.get('aws')
14+
const { credentials, region } = awsCredentials
15+
if (!credentials || !credentials.accessKeyId || !credentials.secretAccessKey || !region) {
1616
console.warn('AWS credentials not configured')
1717
return []
1818
}
1919

20-
const cacheKey = `${region}-${accessKeyId}`
20+
const cacheKey = `${region}-${credentials.accessKeyId}`
2121
const cachedData = this.modelCache[cacheKey]
2222

2323
if (
@@ -31,7 +31,7 @@ export class ModelService {
3131
try {
3232
const models = getModelsForRegion(region as BedrockSupportRegion)
3333

34-
const accountId = await getAccountId(credentials)
34+
const accountId = await getAccountId(awsCredentials)
3535
const promptRouterModels = accountId ? getDefaultPromptRouter(accountId, region) : []
3636
const result = [...models, ...promptRouterModels]
3737
this.modelCache[cacheKey] = [...result, { _timestamp: Date.now() } as any]

‎src/main/api/bedrock/types.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { Message } from '@aws-sdk/client-bedrock-runtime'
2+
import type { AwsCredentialIdentity } from "@smithy/types";
23

34
export type CallConverseAPIProps = {
45
modelId: string
@@ -8,9 +9,8 @@ export type CallConverseAPIProps = {
89
}
910

1011
export type AWSCredentials = {
11-
region: string
12-
accessKeyId: string
13-
secretAccessKey: string
12+
credentials: AwsCredentialIdentity;
13+
region: string;
1414
}
1515

1616
export type InferenceParams = {

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

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -50,16 +50,9 @@ export function getAlternateRegionOnThrottling(
5050
return availableRegions[randomIndex]
5151
}
5252

53-
export async function getAccountId(credentials: AWSCredentials) {
53+
export async function getAccountId(awsCredentials: AWSCredentials) {
5454
try {
55-
const { region, accessKeyId, secretAccessKey } = credentials
56-
const sts = new STSClient({
57-
credentials: {
58-
accessKeyId,
59-
secretAccessKey
60-
},
61-
region
62-
})
55+
const sts = new STSClient(awsCredentials)
6356
const command = new GetCallerIdentityCommand({})
6457
const res = await sts.send(command)
6558
return res.Account

‎src/renderer/src/contexts/SettingsContext.tsx‎

Lines changed: 42 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
PRODUCT_DESIGNER_SYSTEM_PROMPT
1313
} from '@renderer/pages/ChatPage/constants/DEFAULT_AGENTS'
1414
import { InferenceParameters, LLM, BEDROCK_SUPPORTED_REGIONS } from '@/types/llm'
15+
import type { AwsCredentialIdentity } from "@smithy/types";
1516

1617
const DEFAULT_INFERENCE_PARAMS: InferenceParameters = {
1718
maxTokens: 4096,
@@ -141,6 +142,8 @@ export interface SettingsContextType {
141142
setAwsAccessKeyId: (accessKeyId: string) => void
142143
awsSecretAccessKey: string
143144
setAwsSecretAccessKey: (secretAccessKey: string) => void
145+
awsSessionToken: string
146+
setAwsSessionToken: (sessionToken: string) => void
144147

145148
// Custom Agents Settings
146149
customAgents: CustomAgent[]
@@ -215,6 +218,7 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
215218
const [awsRegion, setStateAwsRegion] = useState<string>('')
216219
const [awsAccessKeyId, setStateAwsAccessKeyId] = useState<string>('')
217220
const [awsSecretAccessKey, setStateAwsSecretAccessKey] = useState<string>('')
221+
const [awsSessionToken, setStateAwsSessionToken] = useState<string>('')
218222

219223
// Custom Agents Settings
220224
const [customAgents, setCustomAgents] = useState<CustomAgent[]>([])
@@ -290,6 +294,7 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
290294
setStateAwsRegion(awsConfig.region || '')
291295
setStateAwsAccessKeyId(awsConfig.accessKeyId || '')
292296
setStateAwsSecretAccessKey(awsConfig.secretAccessKey || '')
297+
setStateAwsSessionToken(awsConfig.sessionToken || '')
293298
}
294299

295300
// Load Custom Agents
@@ -442,7 +447,12 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
442447

443448
const setAwsRegion = (region: string) => {
444449
setStateAwsRegion(region)
445-
saveAwsConfig(region, awsAccessKeyId, awsSecretAccessKey)
450+
const credentials: AwsCredentialIdentity = {
451+
accessKeyId: awsAccessKeyId,
452+
secretAccessKey: awsSecretAccessKey,
453+
sessionToken: !awsSessionToken ? undefined : awsSessionToken
454+
}
455+
saveAwsConfig(credentials, region)
446456

447457
// availableFailoverRegions をリセット
448458
setBedrockSettings({
@@ -457,20 +467,36 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
457467

458468
const setAwsAccessKeyId = (accessKeyId: string) => {
459469
setStateAwsAccessKeyId(accessKeyId)
460-
saveAwsConfig(awsRegion, accessKeyId, awsSecretAccessKey)
470+
const credentials: AwsCredentialIdentity = {
471+
accessKeyId,
472+
secretAccessKey: awsSecretAccessKey,
473+
sessionToken: !awsSessionToken ? undefined : awsSessionToken
474+
}
475+
saveAwsConfig(credentials, awsRegion)
461476
}
462477

463478
const setAwsSecretAccessKey = (secretAccessKey: string) => {
464479
setStateAwsSecretAccessKey(secretAccessKey)
465-
saveAwsConfig(awsRegion, awsAccessKeyId, secretAccessKey)
480+
const credentials: AwsCredentialIdentity = {
481+
accessKeyId: awsAccessKeyId,
482+
secretAccessKey,
483+
sessionToken: !awsSessionToken ? undefined : awsSessionToken
484+
}
485+
saveAwsConfig(credentials, awsRegion)
466486
}
467487

468-
const saveAwsConfig = (region: string, accessKeyId: string, secretAccessKey: string) => {
469-
window.store.set('aws', {
470-
region,
471-
accessKeyId,
472-
secretAccessKey
473-
})
488+
const setAwsSessionToken = (sessionToken: string) => {
489+
setStateAwsSessionToken(sessionToken)
490+
const credentials: AwsCredentialIdentity = {
491+
accessKeyId: awsAccessKeyId,
492+
secretAccessKey: awsSecretAccessKey,
493+
sessionToken: !sessionToken ? undefined : sessionToken
494+
}
495+
saveAwsConfig(credentials, awsRegion)
496+
}
497+
498+
const saveAwsConfig = (credentials: AwsCredentialIdentity, region: string) => {
499+
window.store.set('aws', { credentials, region })
474500
}
475501

476502
const saveCustomAgents = (agents: CustomAgent[]) => {
@@ -581,11 +607,11 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
581607
const currentAgent = allAgents.find((a) => a.id === selectedAgentId)
582608
const systemPrompt = currentAgent?.system
583609
? replacePlaceholders(currentAgent?.system, {
584-
projectPath,
585-
allowedCommands: allowedCommands,
586-
knowledgeBases: knowledgeBases,
587-
bedrockAgents: bedrockAgents
588-
})
610+
projectPath,
611+
allowedCommands: allowedCommands,
612+
knowledgeBases: knowledgeBases,
613+
bedrockAgents: bedrockAgents
614+
})
589615
: ''
590616

591617
const setTools = (newTools: ToolState[]) => {
@@ -678,6 +704,8 @@ export const SettingsProvider: React.FC<{ children: React.ReactNode }> = ({ chil
678704
setAwsAccessKeyId,
679705
awsSecretAccessKey,
680706
setAwsSecretAccessKey,
707+
awsSessionToken,
708+
setAwsSessionToken,
681709

682710
// Custom Agents Settings
683711
customAgents,

‎src/renderer/src/pages/SettingPage/SettingPage.tsx‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ export const SettingPage: React.FC = () => {
3030
setAwsAccessKeyId,
3131
awsSecretAccessKey,
3232
setAwsSecretAccessKey,
33+
awsSessionToken,
34+
setAwsSessionToken,
3335
inferenceParams,
3436
updateInferenceParams,
3537
bedrockSettings,
@@ -72,9 +74,11 @@ export const SettingPage: React.FC = () => {
7274
awsRegion={awsRegion}
7375
awsAccessKeyId={awsAccessKeyId}
7476
awsSecretAccessKey={awsSecretAccessKey}
77+
awsSessionToken={awsSessionToken}
7578
onUpdateRegion={setAwsRegion}
7679
onUpdateAccessKeyId={setAwsAccessKeyId}
7780
onUpdateSecretAccessKey={setAwsSecretAccessKey}
81+
onUpdateSessionToken={setAwsSessionToken}
7882
currentLLM={currentLLM}
7983
availableModels={availableModels}
8084
inferenceParams={inferenceParams}

‎src/renderer/src/pages/SettingPage/components/sections/AWSSection.tsx‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,11 @@ interface AWSSectionProps {
1414
awsRegion: string
1515
awsAccessKeyId: string
1616
awsSecretAccessKey: string
17+
awsSessionToken: string
1718
onUpdateRegion: (region: string) => void
1819
onUpdateAccessKeyId: (id: string) => void
1920
onUpdateSecretAccessKey: (key: string) => void
21+
onUpdateSessionToken: (token: string) => void
2022

2123
// Bedrock Settings
2224
currentLLM: LLM
@@ -40,9 +42,11 @@ export const AWSSection: React.FC<AWSSectionProps> = ({
4042
awsRegion,
4143
awsAccessKeyId,
4244
awsSecretAccessKey,
45+
awsSessionToken,
4346
onUpdateRegion,
4447
onUpdateAccessKeyId,
4548
onUpdateSecretAccessKey,
49+
onUpdateSessionToken,
4650

4751
// Bedrock Settings
4852
currentLLM,
@@ -142,6 +146,14 @@ export const AWSSection: React.FC<AWSSectionProps> = ({
142146
onChange={(e) => onUpdateSecretAccessKey(e.target.value)}
143147
/>
144148

149+
<SettingInput
150+
label={t('AWS Session Token (optional)')}
151+
type="password"
152+
placeholder="****************************************"
153+
value={awsSessionToken}
154+
onChange={(e) => onUpdateSessionToken(e.target.value)}
155+
/>
156+
145157
<SettingSelect
146158
label={t('AWS Region')}
147159
value={awsRegion}

0 commit comments

Comments
 (0)