Skip to content

Commit 5b444ef

Browse files
aidandaly24aidandaly24
committed
feat: smart credential detection for multi-agent projects
- Add resolveCredentialStrategy() to detect same vs different API keys - Same key reuses project-scoped credential, different key creates agent-scoped - Clean up agent-scoped credentials on agent removal - Fix deploy flow: always update credentials on deploy (not just manual entry) - Add getAllCredentials() for credential prompt, updateApiKeyProvider() for updates - Fix double-execution bugs in useCdkPreflight and CredentialSourcePrompt - Add unit and integration tests
1 parent 93f6309 commit 5b444ef

21 files changed

Lines changed: 1197 additions & 108 deletions

‎src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2606,7 +2606,7 @@ exports[`Assets Directory Snapshots > Python framework assets > python/python/st
26062606
"import os
26072607
from typing import Optional
26082608
2609-
from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig, RetrievalConfig
2609+
from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig{{#if memoryProviders.[0].strategies.length}}, RetrievalConfig{{/if}}
26102610
from bedrock_agentcore.memory.integrations.strands.session_manager import AgentCoreMemorySessionManager
26112611
26122612
MEMORY_ID = os.getenv("{{memoryProviders.[0].envVarName}}")
@@ -2616,16 +2616,28 @@ def get_memory_session_manager(session_id: str, actor_id: str) -> Optional[Agent
26162616
if not MEMORY_ID:
26172617
return None
26182618
2619+
{{#if memoryProviders.[0].strategies.length}}
2620+
retrieval_config = {
2621+
{{#if (includes memoryProviders.[0].strategies "SEMANTIC")}}
2622+
f"/users/{actor_id}/facts": RetrievalConfig(top_k=3, relevance_score=0.5),
2623+
{{/if}}
2624+
{{#if (includes memoryProviders.[0].strategies "USER_PREFERENCE")}}
2625+
f"/users/{actor_id}/preferences": RetrievalConfig(top_k=3, relevance_score=0.5),
2626+
{{/if}}
2627+
{{#if (includes memoryProviders.[0].strategies "SUMMARIZATION")}}
2628+
f"/summaries/{actor_id}/{session_id}": RetrievalConfig(top_k=3, relevance_score=0.5),
2629+
{{/if}}
2630+
}
2631+
{{/if}}
2632+
26192633
return AgentCoreMemorySessionManager(
26202634
AgentCoreMemoryConfig(
26212635
memory_id=MEMORY_ID,
26222636
session_id=session_id,
26232637
actor_id=actor_id,
2624-
retrieval_config={
2625-
f"/users/{actor_id}/facts": RetrievalConfig(top_k=3, relevance_score=0.5),
2626-
f"/users/{actor_id}/preferences": RetrievalConfig(top_k=3, relevance_score=0.5),
2627-
f"/summaries/{actor_id}/{session_id}": RetrievalConfig(top_k=3, relevance_score=0.5),
2628-
}
2638+
{{#if memoryProviders.[0].strategies.length}}
2639+
retrieval_config=retrieval_config,
2640+
{{/if}}
26292641
),
26302642
REGION
26312643
)

‎src/cli/commands/add/__tests__/multi-agent-credentials.test.ts‎

Lines changed: 281 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,281 @@
1+
/* eslint-disable security/detect-non-literal-fs-filename */
2+
import { runCLI } from '../../../../test-utils/index.js';
3+
import { randomUUID } from 'node:crypto';
4+
import { mkdir, readFile, rm, writeFile } from 'node:fs/promises';
5+
import { tmpdir } from 'node:os';
6+
import { join } from 'node:path';
7+
import { afterAll, beforeAll, describe, expect, it } from 'vitest';
8+
9+
/**
10+
* Integration tests for multi-agent credential behavior (Option C).
11+
*
12+
* Tests the smart credential detection:
13+
* - Same API key → reuse existing project-scoped credential
14+
* - Different API key → create agent-scoped credential
15+
* - Remove agent → clean up agent-scoped credentials
16+
*/
17+
describe('multi-agent credential behavior', () => {
18+
let testDir: string;
19+
let projectDir: string;
20+
const projectName = 'MultiAgentProj';
21+
22+
beforeAll(async () => {
23+
testDir = join(tmpdir(), `agentcore-multi-agent-cred-${randomUUID()}`);
24+
await mkdir(testDir, { recursive: true });
25+
26+
// Create project without agent
27+
const result = await runCLI(['create', '--name', projectName, '--no-agent'], testDir);
28+
if (result.exitCode !== 0) {
29+
throw new Error(`Failed to create project: ${result.stdout} ${result.stderr}`);
30+
}
31+
projectDir = join(testDir, projectName);
32+
});
33+
34+
afterAll(async () => {
35+
await rm(testDir, { recursive: true, force: true });
36+
});
37+
38+
async function readProjectSpec() {
39+
const content = await readFile(join(projectDir, 'agentcore/agentcore.json'), 'utf-8');
40+
return JSON.parse(content);
41+
}
42+
43+
async function readEnvLocal() {
44+
try {
45+
return await readFile(join(projectDir, 'agentcore/.env.local'), 'utf-8');
46+
} catch {
47+
return '';
48+
}
49+
}
50+
51+
describe('credential reuse with same API key', () => {
52+
it('first agent creates project-scoped credential', async () => {
53+
const result = await runCLI(
54+
[
55+
'add',
56+
'agent',
57+
'--name',
58+
'Agent1',
59+
'--language',
60+
'Python',
61+
'--framework',
62+
'Strands',
63+
'--model-provider',
64+
'Gemini',
65+
'--api-key',
66+
'KEY1',
67+
'--memory',
68+
'none',
69+
'--json',
70+
],
71+
projectDir
72+
);
73+
74+
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
75+
76+
const spec = await readProjectSpec();
77+
expect(spec.credentials).toHaveLength(1);
78+
expect(spec.credentials[0].name).toBe(`${projectName}Gemini`);
79+
80+
const env = await readEnvLocal();
81+
expect(env).toContain('AGENTCORE_CREDENTIAL_MULTIAGENTPROJGEMINI=');
82+
expect(env).toContain('KEY1');
83+
});
84+
85+
it('second agent with same key reuses credential (no duplicate)', async () => {
86+
const result = await runCLI(
87+
[
88+
'add',
89+
'agent',
90+
'--name',
91+
'Agent2',
92+
'--language',
93+
'Python',
94+
'--framework',
95+
'Strands',
96+
'--model-provider',
97+
'Gemini',
98+
'--api-key',
99+
'KEY1',
100+
'--memory',
101+
'none',
102+
'--json',
103+
],
104+
projectDir
105+
);
106+
107+
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
108+
109+
const spec = await readProjectSpec();
110+
// Should still have only 1 credential (reused)
111+
expect(spec.credentials).toHaveLength(1);
112+
expect(spec.credentials[0].name).toBe(`${projectName}Gemini`);
113+
114+
// Should have 2 agents
115+
expect(spec.agents).toHaveLength(2);
116+
});
117+
});
118+
119+
describe('agent-scoped credential with different API key', () => {
120+
it('third agent with different key creates agent-scoped credential', async () => {
121+
const result = await runCLI(
122+
[
123+
'add',
124+
'agent',
125+
'--name',
126+
'Agent3',
127+
'--language',
128+
'Python',
129+
'--framework',
130+
'Strands',
131+
'--model-provider',
132+
'Gemini',
133+
'--api-key',
134+
'KEY2',
135+
'--memory',
136+
'none',
137+
'--json',
138+
],
139+
projectDir
140+
);
141+
142+
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
143+
144+
const spec = await readProjectSpec();
145+
// Should now have 2 credentials
146+
expect(spec.credentials).toHaveLength(2);
147+
148+
const credNames = spec.credentials.map((c: { name: string }) => c.name);
149+
expect(credNames).toContain(`${projectName}Gemini`);
150+
expect(credNames).toContain(`${projectName}Agent3Gemini`);
151+
152+
// Should have 3 agents
153+
expect(spec.agents).toHaveLength(3);
154+
155+
// .env.local should have both keys
156+
const env = await readEnvLocal();
157+
expect(env).toContain('AGENTCORE_CREDENTIAL_MULTIAGENTPROJGEMINI=');
158+
expect(env).toContain('KEY1');
159+
expect(env).toContain('AGENTCORE_CREDENTIAL_MULTIAGENTPROJAGENT3GEMINI=');
160+
expect(env).toContain('KEY2');
161+
162+
// Generated code should reference correct credentials
163+
const agent1Code = await readFile(join(projectDir, 'app/Agent1/model/load.py'), 'utf-8');
164+
expect(agent1Code).toContain(`IDENTITY_PROVIDER_NAME = "${projectName}Gemini"`);
165+
166+
const agent3Code = await readFile(join(projectDir, 'app/Agent3/model/load.py'), 'utf-8');
167+
expect(agent3Code).toContain(`IDENTITY_PROVIDER_NAME = "${projectName}Agent3Gemini"`);
168+
});
169+
});
170+
171+
describe('credential cleanup on agent removal', () => {
172+
it('removing agent with agent-scoped credential cleans up credential', async () => {
173+
const result = await runCLI(['remove', 'agent', '--name', 'Agent3', '--json'], projectDir);
174+
175+
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
176+
177+
const spec = await readProjectSpec();
178+
// Should be back to 1 credential (agent-scoped removed)
179+
expect(spec.credentials).toHaveLength(1);
180+
expect(spec.credentials[0].name).toBe(`${projectName}Gemini`);
181+
182+
// Should have 2 agents
183+
expect(spec.agents).toHaveLength(2);
184+
});
185+
186+
it('removing agent with shared credential does NOT remove credential', async () => {
187+
// Remove Agent2 (uses shared project-scoped credential)
188+
const result = await runCLI(['remove', 'agent', '--name', 'Agent2', '--json'], projectDir);
189+
190+
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
191+
192+
const spec = await readProjectSpec();
193+
// Credential should still exist (shared with Agent1)
194+
expect(spec.credentials).toHaveLength(1);
195+
expect(spec.credentials[0].name).toBe(`${projectName}Gemini`);
196+
197+
// Should have 1 agent
198+
expect(spec.agents).toHaveLength(1);
199+
});
200+
});
201+
202+
describe('BYO (bring-your-own) agent path', () => {
203+
it('BYO agent with same key reuses credential', async () => {
204+
// Create a code directory for BYO agent
205+
const byoDir = join(projectDir, 'app/ByoAgent');
206+
await mkdir(byoDir, { recursive: true });
207+
await writeFile(join(byoDir, 'main.py'), '# BYO agent');
208+
209+
const result = await runCLI(
210+
[
211+
'add',
212+
'agent',
213+
'--name',
214+
'ByoAgent',
215+
'--type',
216+
'byo',
217+
'--language',
218+
'Python',
219+
'--framework',
220+
'Strands',
221+
'--code-location',
222+
'app/ByoAgent/',
223+
'--model-provider',
224+
'Gemini',
225+
'--api-key',
226+
'KEY1',
227+
'--json',
228+
],
229+
projectDir
230+
);
231+
232+
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
233+
234+
const spec = await readProjectSpec();
235+
// Should still have only 1 credential (reused)
236+
expect(spec.credentials).toHaveLength(1);
237+
expect(spec.credentials[0].name).toBe(`${projectName}Gemini`);
238+
239+
// Should have 2 agents now
240+
expect(spec.agents).toHaveLength(2);
241+
});
242+
243+
it('BYO agent with different key creates agent-scoped credential', async () => {
244+
const byoDir2 = join(projectDir, 'app/ByoAgent2');
245+
await mkdir(byoDir2, { recursive: true });
246+
await writeFile(join(byoDir2, 'main.py'), '# BYO agent 2');
247+
248+
const result = await runCLI(
249+
[
250+
'add',
251+
'agent',
252+
'--name',
253+
'ByoAgent2',
254+
'--type',
255+
'byo',
256+
'--language',
257+
'Python',
258+
'--framework',
259+
'Strands',
260+
'--code-location',
261+
'app/ByoAgent2/',
262+
'--model-provider',
263+
'Gemini',
264+
'--api-key',
265+
'DIFFERENT_KEY',
266+
'--json',
267+
],
268+
projectDir
269+
);
270+
271+
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
272+
273+
const spec = await readProjectSpec();
274+
// Should have 2 credentials now
275+
expect(spec.credentials).toHaveLength(2);
276+
const credNames = spec.credentials.map((c: { name: string }) => c.name);
277+
expect(credNames).toContain(`${projectName}Gemini`);
278+
expect(credNames).toContain(`${projectName}ByoAgent2Gemini`);
279+
});
280+
});
281+
});

0 commit comments

Comments
 (0)