Skip to content

Commit 959b709

Browse files
aidandaly24aidandaly24
committed
fix: reuse credentials across agents with same key, preserve on removal
- Check ALL existing credentials for matching API key (not just project-scoped) - Agent3 can now reuse Agent2's credential if they share the same key - Preserve credentials on agent removal for potential reuse - If agent re-added with different key, credential is updated
1 parent 5c07843 commit 959b709

4 files changed

Lines changed: 132 additions & 54 deletions

File tree

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

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -168,31 +168,31 @@ describe('multi-agent credential behavior', () => {
168168
});
169169
});
170170

171-
describe('credential cleanup on agent removal', () => {
172-
it('removing agent with agent-scoped credential cleans up credential', async () => {
171+
describe('credential persistence on agent removal', () => {
172+
it('removing agent preserves agent-scoped credential for reuse', async () => {
173173
const result = await runCLI(['remove', 'agent', '--name', 'Agent3', '--json'], projectDir);
174174

175175
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
176176

177177
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`);
178+
// Credentials preserved (both project-scoped and agent-scoped)
179+
expect(spec.credentials).toHaveLength(2);
180+
expect(spec.credentials.map(c => c.name)).toContain(`${projectName}Gemini`);
181+
expect(spec.credentials.map(c => c.name)).toContain(`${projectName}Agent3Gemini`);
181182

182183
// Should have 2 agents
183184
expect(spec.agents).toHaveLength(2);
184185
});
185186

186-
it('removing agent with shared credential does NOT remove credential', async () => {
187+
it('removing agent with shared credential preserves credential', async () => {
187188
// Remove Agent2 (uses shared project-scoped credential)
188189
const result = await runCLI(['remove', 'agent', '--name', 'Agent2', '--json'], projectDir);
189190

190191
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
191192

192193
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`);
194+
// Both credentials still exist
195+
expect(spec.credentials).toHaveLength(2);
196196

197197
// Should have 1 agent
198198
expect(spec.agents).toHaveLength(1);
@@ -206,6 +206,9 @@ describe('multi-agent credential behavior', () => {
206206
await mkdir(byoDir, { recursive: true });
207207
await writeFile(join(byoDir, 'main.py'), '# BYO agent');
208208

209+
const specBefore = await readProjectSpec();
210+
const credCountBefore = specBefore.credentials.length;
211+
209212
const result = await runCLI(
210213
[
211214
'add',
@@ -232,19 +235,18 @@ describe('multi-agent credential behavior', () => {
232235
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
233236

234237
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);
238+
// Should still have same number of credentials (reused)
239+
expect(spec.credentials).toHaveLength(credCountBefore);
241240
});
242241

243242
it('BYO agent with different key creates agent-scoped credential', async () => {
244243
const byoDir2 = join(projectDir, 'app/ByoAgent2');
245244
await mkdir(byoDir2, { recursive: true });
246245
await writeFile(join(byoDir2, 'main.py'), '# BYO agent 2');
247246

247+
const specBefore = await readProjectSpec();
248+
const credCountBefore = specBefore.credentials.length;
249+
248250
const result = await runCLI(
249251
[
250252
'add',
@@ -271,10 +273,9 @@ describe('multi-agent credential behavior', () => {
271273
expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
272274

273275
const spec = await readProjectSpec();
274-
// Should have 2 credentials now
275-
expect(spec.credentials).toHaveLength(2);
276+
// Should have one more credential
277+
expect(spec.credentials).toHaveLength(credCountBefore + 1);
276278
const credNames = spec.credentials.map((c: { name: string }) => c.name);
277-
expect(credNames).toContain(`${projectName}Gemini`);
278279
expect(credNames).toContain(`${projectName}ByoAgent2Gemini`);
279280
});
280281
});

‎src/cli/operations/identity/__tests__/resolve-credential-strategy.test.ts‎

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ describe('resolveCredentialStrategy', () => {
138138
});
139139
});
140140

141-
it('creates project-scoped credential when existing key is undefined', async () => {
141+
it('creates agent-scoped credential when no existing keys can be read', async () => {
142142
mockGetEnvVar.mockResolvedValue(undefined);
143143

144144
const result = await resolveCredentialStrategy(
@@ -150,13 +150,105 @@ describe('resolveCredentialStrategy', () => {
150150
existingCredentials
151151
);
152152

153+
// Can't verify existing key matches, so create agent-scoped to be safe
153154
expect(result).toEqual({
154155
reuse: false,
156+
credentialName: 'MyProjectAgent1Gemini',
157+
envVarName: 'AGENTCORE_CREDENTIAL_MYPROJECTAGENT1GEMINI',
158+
isAgentScoped: true,
159+
});
160+
});
161+
});
162+
163+
describe('reuses agent-scoped credential when keys match', () => {
164+
it('agent3 reuses agent2 credential when both use same secondary key', async () => {
165+
// Scenario: agent1 uses mainKey (project-scoped), agent2 uses secondaryKey (agent-scoped)
166+
// agent3 also uses secondaryKey - should reuse agent2's credential
167+
const existingCredentials: Credential[] = [
168+
{ name: 'MyProjectGemini', type: 'ApiKeyCredentialProvider' },
169+
{ name: 'MyProjectAgent2Gemini', type: 'ApiKeyCredentialProvider' },
170+
];
171+
172+
mockGetEnvVar.mockImplementation((envVar: string) => {
173+
if (envVar === 'AGENTCORE_CREDENTIAL_MYPROJECTGEMINI') return Promise.resolve('main-key');
174+
if (envVar === 'AGENTCORE_CREDENTIAL_MYPROJECTAGENT2GEMINI') return Promise.resolve('secondary-key');
175+
return Promise.resolve(undefined);
176+
});
177+
178+
const result = await resolveCredentialStrategy(
179+
projectName,
180+
'Agent3',
181+
'Gemini',
182+
'secondary-key',
183+
configBaseDir,
184+
existingCredentials
185+
);
186+
187+
expect(result).toEqual({
188+
reuse: true,
189+
credentialName: 'MyProjectAgent2Gemini',
190+
envVarName: 'AGENTCORE_CREDENTIAL_MYPROJECTAGENT2GEMINI',
191+
isAgentScoped: true,
192+
});
193+
});
194+
195+
it('agent3 reuses project-scoped credential when using main key', async () => {
196+
const existingCredentials: Credential[] = [
197+
{ name: 'MyProjectGemini', type: 'ApiKeyCredentialProvider' },
198+
{ name: 'MyProjectAgent2Gemini', type: 'ApiKeyCredentialProvider' },
199+
];
200+
201+
mockGetEnvVar.mockImplementation((envVar: string) => {
202+
if (envVar === 'AGENTCORE_CREDENTIAL_MYPROJECTGEMINI') return Promise.resolve('main-key');
203+
if (envVar === 'AGENTCORE_CREDENTIAL_MYPROJECTAGENT2GEMINI') return Promise.resolve('secondary-key');
204+
return Promise.resolve(undefined);
205+
});
206+
207+
const result = await resolveCredentialStrategy(
208+
projectName,
209+
'Agent3',
210+
'Gemini',
211+
'main-key',
212+
configBaseDir,
213+
existingCredentials
214+
);
215+
216+
expect(result).toEqual({
217+
reuse: true,
155218
credentialName: 'MyProjectGemini',
156219
envVarName: 'AGENTCORE_CREDENTIAL_MYPROJECTGEMINI',
157220
isAgentScoped: false,
158221
});
159222
});
223+
224+
it('creates new agent-scoped credential when key matches no existing credential', async () => {
225+
const existingCredentials: Credential[] = [
226+
{ name: 'MyProjectGemini', type: 'ApiKeyCredentialProvider' },
227+
{ name: 'MyProjectAgent2Gemini', type: 'ApiKeyCredentialProvider' },
228+
];
229+
230+
mockGetEnvVar.mockImplementation((envVar: string) => {
231+
if (envVar === 'AGENTCORE_CREDENTIAL_MYPROJECTGEMINI') return Promise.resolve('main-key');
232+
if (envVar === 'AGENTCORE_CREDENTIAL_MYPROJECTAGENT2GEMINI') return Promise.resolve('secondary-key');
233+
return Promise.resolve(undefined);
234+
});
235+
236+
const result = await resolveCredentialStrategy(
237+
projectName,
238+
'Agent3',
239+
'Gemini',
240+
'third-key',
241+
configBaseDir,
242+
existingCredentials
243+
);
244+
245+
expect(result).toEqual({
246+
reuse: false,
247+
credentialName: 'MyProjectAgent3Gemini',
248+
envVarName: 'AGENTCORE_CREDENTIAL_MYPROJECTAGENT3GEMINI',
249+
isAgentScoped: true,
250+
});
251+
});
160252
});
161253

162254
describe('credential name format', () => {

‎src/cli/operations/identity/create-identity.ts‎

Lines changed: 17 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,8 @@ export function computeDefaultCredentialEnvVarName(credentialName: string): stri
3838
* - Bedrock uses IAM, no credential needed
3939
* - No API key provided, no credential needed
4040
* - No existing credential for provider → create project-scoped
41-
* - Existing credential but can't read key → create project-scoped (treat as fresh)
42-
* - Existing credential with matching key → reuse
43-
* - Existing credential with different key → create agent-scoped
41+
* - Any existing credential with matching key → reuse it
42+
* - No matching key → create agent-scoped (or project-scoped if first)
4443
*/
4544
export async function resolveCredentialStrategy(
4645
projectName: string,
@@ -60,31 +59,27 @@ export async function resolveCredentialStrategy(
6059
return { reuse: true, credentialName: '', envVarName: '', isAgentScoped: false };
6160
}
6261

63-
// Check for existing project-scoped credential
62+
// Check ALL existing credentials for a matching API key
63+
for (const cred of existingCredentials) {
64+
const envVarName = computeDefaultCredentialEnvVarName(cred.name);
65+
const existingApiKey = await getEnvVar(envVarName, configBaseDir);
66+
if (existingApiKey === newApiKey) {
67+
const isAgentScoped = cred.name !== `${projectName}${modelProvider}`;
68+
return { reuse: true, credentialName: cred.name, envVarName, isAgentScoped };
69+
}
70+
}
71+
72+
// No matching key found - create new credential
6473
const projectScopedName = `${projectName}${modelProvider}`;
65-
const existingCredential = existingCredentials.find(c => c.name === projectScopedName);
74+
const hasProjectScoped = existingCredentials.some(c => c.name === projectScopedName);
6675

67-
if (!existingCredential) {
68-
// First agent with this provider - create project-scoped credential
76+
if (!hasProjectScoped) {
77+
// First agent with this provider - create project-scoped
6978
const envVarName = computeDefaultCredentialEnvVarName(projectScopedName);
7079
return { reuse: false, credentialName: projectScopedName, envVarName, isAgentScoped: false };
7180
}
7281

73-
// Credential exists - compare API keys
74-
const existingEnvVarName = computeDefaultCredentialEnvVarName(projectScopedName);
75-
const existingApiKey = await getEnvVar(existingEnvVarName, configBaseDir);
76-
77-
if (existingApiKey === undefined) {
78-
// Can't read existing key - treat as no existing credential
79-
return { reuse: false, credentialName: projectScopedName, envVarName: existingEnvVarName, isAgentScoped: false };
80-
}
81-
82-
if (existingApiKey === newApiKey) {
83-
// Same key - reuse existing credential
84-
return { reuse: true, credentialName: projectScopedName, envVarName: existingEnvVarName, isAgentScoped: false };
85-
}
86-
87-
// Different key - create agent-scoped credential
82+
// Project-scoped exists with different key - create agent-scoped
8883
const agentScopedName = `${projectName}${agentName}${modelProvider}`;
8984
const agentScopedEnvVarName = computeDefaultCredentialEnvVarName(agentScopedName);
9085
return { reuse: false, credentialName: agentScopedName, envVarName: agentScopedEnvVarName, isAgentScoped: true };

‎src/cli/operations/remove/remove-agent.ts‎

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ export async function getRemovableAgents(): Promise<string[]> {
3737

3838
/**
3939
* Preview what will be removed when removing an agent.
40+
* Note: Credentials are preserved to allow reuse if agent is re-added.
4041
*/
4142
export async function previewRemoveAgent(agentName: string): Promise<RemovalPreview> {
4243
const configIO = new ConfigIO();
@@ -50,17 +51,9 @@ export async function previewRemoveAgent(agentName: string): Promise<RemovalPrev
5051
const summary: string[] = [`Removing agent: ${agentName}`];
5152
const schemaChanges: SchemaChange[] = [];
5253

53-
// Identify agent-scoped credentials
54-
const agentCredentials = getAgentScopedCredentials(project.name, agentName, project.credentials);
55-
if (agentCredentials.length > 0) {
56-
summary.push(`Will remove ${agentCredentials.length} agent-scoped credential(s):`);
57-
agentCredentials.forEach(c => summary.push(` - ${c.name}`));
58-
}
59-
6054
const afterSpec = {
6155
...project,
6256
agents: project.agents.filter(a => a.name !== agentName),
63-
credentials: project.credentials.filter(c => !agentCredentials.some(ac => ac.name === c.name)),
6457
};
6558

6659
schemaChanges.push({
@@ -74,6 +67,7 @@ export async function previewRemoveAgent(agentName: string): Promise<RemovalPrev
7467

7568
/**
7669
* Remove an agent from the project.
70+
* Note: Credentials are preserved to allow reuse if agent is re-added.
7771
*/
7872
export async function removeAgent(agentName: string): Promise<RemovalResult> {
7973
try {
@@ -85,13 +79,9 @@ export async function removeAgent(agentName: string): Promise<RemovalResult> {
8579
return { ok: false, error: `Agent "${agentName}" not found.` };
8680
}
8781

88-
// Remove agent
82+
// Remove agent (credentials preserved for potential reuse)
8983
project.agents.splice(agentIndex, 1);
9084

91-
// Remove agent-scoped credentials
92-
const agentCredentials = getAgentScopedCredentials(project.name, agentName, project.credentials);
93-
project.credentials = project.credentials.filter(c => !agentCredentials.some(ac => ac.name === c.name));
94-
9585
await configIO.writeProjectSpec(project);
9686

9787
return { ok: true };

0 commit comments

Comments
 (0)