Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 74 additions & 2 deletions src/cli/operations/mcp/__tests__/create-mcp.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { SKIP_FOR_NOW } from '../../../tui/screens/mcp/types.js';
import type { AddGatewayTargetConfig } from '../../../tui/screens/mcp/types.js';
import { createExternalGatewayTarget } from '../create-mcp.js';
import type { AddGatewayConfig, AddGatewayTargetConfig } from '../../../tui/screens/mcp/types.js';
import { createExternalGatewayTarget, createGatewayFromWizard, getUnassignedTargets } from '../create-mcp.js';
import { afterEach, describe, expect, it, vi } from 'vitest';

const { mockReadMcpSpec, mockWriteMcpSpec, mockConfigExists, mockReadProjectSpec } = vi.hoisted(() => ({
Expand Down Expand Up @@ -131,3 +131,75 @@ describe('createExternalGatewayTarget', () => {
expect(target.outboundAuth).toEqual({ type: 'API_KEY', credentialName: 'my-cred' });
});
});

describe('getUnassignedTargets', () => {
afterEach(() => vi.clearAllMocks());

it('returns unassigned targets from mcp spec', async () => {
mockConfigExists.mockReturnValue(true);
mockReadMcpSpec.mockResolvedValue({
agentCoreGateways: [],
unassignedTargets: [{ name: 't1' }, { name: 't2' }],
});

const result = await getUnassignedTargets();
expect(result).toHaveLength(2);
expect(result[0]!.name).toBe('t1');
});

it('returns empty array when no mcp config exists', async () => {
mockConfigExists.mockReturnValue(false);
expect(await getUnassignedTargets()).toEqual([]);
});

it('returns empty array when unassignedTargets field is missing', async () => {
mockConfigExists.mockReturnValue(true);
mockReadMcpSpec.mockResolvedValue({ agentCoreGateways: [] });
expect(await getUnassignedTargets()).toEqual([]);
});
});

describe('createGatewayFromWizard with selectedTargets', () => {
afterEach(() => vi.clearAllMocks());

function makeGatewayConfig(overrides: Partial<AddGatewayConfig> = {}): AddGatewayConfig {
return {
name: 'new-gateway',
authorizerType: 'AWS_IAM',
...overrides,
} as AddGatewayConfig;
}

it('moves selected targets to new gateway and removes from unassigned', async () => {
mockConfigExists.mockReturnValue(true);
mockReadMcpSpec.mockResolvedValue({
agentCoreGateways: [],
unassignedTargets: [
{ name: 'target-a', targetType: 'mcpServer' },
{ name: 'target-b', targetType: 'mcpServer' },
{ name: 'target-c', targetType: 'mcpServer' },
],
});

await createGatewayFromWizard(makeGatewayConfig({ selectedTargets: ['target-a', 'target-c'] }));

const written = mockWriteMcpSpec.mock.calls[0]![0];
const gateway = written.agentCoreGateways.find((g: { name: string }) => g.name === 'new-gateway');
expect(gateway.targets).toHaveLength(2);
expect(gateway.targets[0]!.name).toBe('target-a');
expect(gateway.targets[1]!.name).toBe('target-c');
expect(written.unassignedTargets).toHaveLength(1);
expect(written.unassignedTargets[0]!.name).toBe('target-b');
});

it('creates gateway with empty targets when no selectedTargets', async () => {
mockConfigExists.mockReturnValue(true);
mockReadMcpSpec.mockResolvedValue({ agentCoreGateways: [] });

await createGatewayFromWizard(makeGatewayConfig());

const written = mockWriteMcpSpec.mock.calls[0]![0];
const gateway = written.agentCoreGateways.find((g: { name: string }) => g.name === 'new-gateway');
expect(gateway.targets).toHaveLength(0);
});
});
37 changes: 36 additions & 1 deletion src/cli/operations/mcp/create-mcp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,22 @@ function buildAuthorizerConfiguration(config: AddGatewayConfig): AgentCoreGatewa
};
}

/**
* Get list of unassigned targets from MCP spec.
*/
export async function getUnassignedTargets(): Promise<AgentCoreGatewayTarget[]> {
try {
const configIO = new ConfigIO();
if (!configIO.configExists('mcp')) {
return [];
}
const mcpSpec = await configIO.readMcpSpec();
return mcpSpec.unassignedTargets ?? [];
} catch {
return [];
}
}

/**
* Get list of existing gateway names from project spec.
*/
Expand Down Expand Up @@ -160,15 +176,34 @@ export async function createGatewayFromWizard(config: AddGatewayConfig): Promise
throw new Error(`Gateway "${config.name}" already exists.`);
}

// Collect selected unassigned targets
const selectedTargets: AgentCoreGatewayTarget[] = [];
if (config.selectedTargets && config.selectedTargets.length > 0) {
const unassignedTargets = mcpSpec.unassignedTargets ?? [];
for (const targetName of config.selectedTargets) {
const target = unassignedTargets.find(t => t.name === targetName);
if (target) {
selectedTargets.push(target);
}
}
}

const gateway: AgentCoreGateway = {
name: config.name,
description: config.description,
targets: [],
targets: selectedTargets,
authorizerType: config.authorizerType,
authorizerConfiguration: buildAuthorizerConfiguration(config),
};

mcpSpec.agentCoreGateways.push(gateway);

// Remove selected targets from unassigned targets
if (config.selectedTargets && config.selectedTargets.length > 0) {
const selected = config.selectedTargets;
mcpSpec.unassignedTargets = (mcpSpec.unassignedTargets ?? []).filter(t => !selected.includes(t.name));
}

await configIO.writeMcpSpec(mcpSpec);

return { name: config.name };
Expand Down
68 changes: 68 additions & 0 deletions src/cli/operations/remove/__tests__/remove-gateway.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
import { previewRemoveGateway, removeGateway } from '../remove-gateway.js';
import { afterEach, describe, expect, it, vi } from 'vitest';

const { mockReadMcpSpec, mockWriteMcpSpec, mockConfigExists } = vi.hoisted(() => ({
mockReadMcpSpec: vi.fn(),
mockWriteMcpSpec: vi.fn(),
mockConfigExists: vi.fn(),
}));

vi.mock('../../../../lib/index.js', () => ({
ConfigIO: class {
configExists = mockConfigExists;
readMcpSpec = mockReadMcpSpec;
writeMcpSpec = mockWriteMcpSpec;
},
}));

describe('removeGateway', () => {
afterEach(() => vi.clearAllMocks());

it('moves gateway targets to unassignedTargets on removal, preserving existing', async () => {
mockReadMcpSpec.mockResolvedValue({
agentCoreGateways: [
{ name: 'gw-to-remove', targets: [{ name: 'target-1' }, { name: 'target-2' }] },
{ name: 'other-gw', targets: [] },
],
unassignedTargets: [{ name: 'already-unassigned' }],
});

const result = await removeGateway('gw-to-remove');

expect(result.ok).toBe(true);
const written = mockWriteMcpSpec.mock.calls[0]![0];
expect(written.agentCoreGateways).toHaveLength(1);
expect(written.agentCoreGateways[0]!.name).toBe('other-gw');
expect(written.unassignedTargets).toHaveLength(3);
expect(written.unassignedTargets[0]!.name).toBe('already-unassigned');
expect(written.unassignedTargets[1]!.name).toBe('target-1');
expect(written.unassignedTargets[2]!.name).toBe('target-2');
});

it('does not modify unassignedTargets when gateway has no targets', async () => {
mockReadMcpSpec.mockResolvedValue({
agentCoreGateways: [{ name: 'empty-gw', targets: [] }],
});

const result = await removeGateway('empty-gw');

expect(result.ok).toBe(true);
const written = mockWriteMcpSpec.mock.calls[0]![0];
expect(written.agentCoreGateways).toHaveLength(0);
expect(written.unassignedTargets).toBeUndefined();
});
});

describe('previewRemoveGateway', () => {
afterEach(() => vi.clearAllMocks());

it('shows "will become unassigned" warning when gateway has targets', async () => {
mockReadMcpSpec.mockResolvedValue({
agentCoreGateways: [{ name: 'my-gw', targets: [{ name: 't1' }, { name: 't2' }] }],
});

const preview = await previewRemoveGateway('my-gw');

expect(preview.summary.some(s => s.includes('2 target(s) will become unassigned'))).toBe(true);
});
});
10 changes: 9 additions & 1 deletion src/cli/operations/remove/remove-gateway.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ export async function previewRemoveGateway(gatewayName: string): Promise<Removal
const schemaChanges: SchemaChange[] = [];

if (gateway.targets.length > 0) {
summary.push(`Note: ${gateway.targets.length} target(s) behind this gateway will become orphaned`);
summary.push(`Note: ${gateway.targets.length} target(s) will become unassigned`);
}

// Compute schema changes
Expand All @@ -52,9 +52,17 @@ export async function previewRemoveGateway(gatewayName: string): Promise<Removal
* Compute the MCP spec after removing a gateway.
*/
function computeRemovedGatewayMcpSpec(mcpSpec: AgentCoreMcpSpec, gatewayName: string): AgentCoreMcpSpec {
const gatewayToRemove = mcpSpec.agentCoreGateways.find(g => g.name === gatewayName);
const targetsToPreserve = gatewayToRemove?.targets ?? [];

return {
...mcpSpec,
agentCoreGateways: mcpSpec.agentCoreGateways.filter(g => g.name !== gatewayName),
// Preserve gateway's targets as unassigned so the user doesn't lose them.
// Only add the field if there are targets to preserve or unassignedTargets already exists.
...(targetsToPreserve.length > 0 || mcpSpec.unassignedTargets
? { unassignedTargets: [...(mcpSpec.unassignedTargets ?? []), ...targetsToPreserve] }
: {}),
};
}

Expand Down
20 changes: 20 additions & 0 deletions src/cli/tui/hooks/useCreateMcp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import {
getAvailableAgents,
getExistingGateways,
getExistingToolNames,
getUnassignedTargets,
} from '../../operations/mcp/create-mcp';
import type { AddGatewayConfig, AddGatewayTargetConfig } from '../screens/mcp/types';
import { useCallback, useEffect, useState } from 'react';
Expand Down Expand Up @@ -117,3 +118,22 @@ export function useExistingToolNames() {

return { toolNames, refresh };
}

export function useUnassignedTargets() {
const [targets, setTargets] = useState<string[]>([]);

useEffect(() => {
async function load() {
const result = await getUnassignedTargets();
setTargets(result.map(t => t.name));
}
void load();
}, []);

const refresh = useCallback(async () => {
const result = await getUnassignedTargets();
setTargets(result.map(t => t.name));
}, []);

return { targets, refresh };
}
5 changes: 4 additions & 1 deletion src/cli/tui/screens/mcp/AddGatewayFlow.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import type { SelectableItem } from '../../components';
import { HELP_TEXT } from '../../constants';
import { useListNavigation } from '../../hooks';
import { useAgents, useAttachGateway, useGateways } from '../../hooks/useAttach';
import { useCreateGateway, useExistingGateways } from '../../hooks/useCreateMcp';
import { useCreateGateway, useExistingGateways, useUnassignedTargets } from '../../hooks/useCreateMcp';
import { AddSuccessScreen } from '../add/AddSuccessScreen';
import { AddGatewayScreen } from './AddGatewayScreen';
import type { AddGatewayConfig } from './types';
Expand Down Expand Up @@ -55,6 +55,7 @@ export function AddGatewayFlow({
}: AddGatewayFlowProps) {
const { createGateway, reset: resetCreate } = useCreateGateway();
const { gateways: existingGateways, refresh: refreshGateways } = useExistingGateways();
const { targets: unassignedTargets } = useUnassignedTargets();
const [flow, setFlow] = useState<FlowState>({ name: 'mode-select' });

// Bind flow hooks
Expand Down Expand Up @@ -157,6 +158,7 @@ export function AddGatewayFlow({
<AddGatewayScreen
existingGateways={existingGateways}
availableAgents={availableAgents}
unassignedTargets={unassignedTargets}
onComplete={handleCreateComplete}
onExit={onBack}
/>
Expand All @@ -183,6 +185,7 @@ export function AddGatewayFlow({
<AddGatewayScreen
existingGateways={existingGateways}
availableAgents={availableAgents}
unassignedTargets={unassignedTargets}
onComplete={handleCreateComplete}
onExit={() => setFlow({ name: 'mode-select' })}
/>
Expand Down
Loading
Loading