feat: assign unassigned targets to gateways and preserve targets on removal (#410)

* feat: assign unassigned targets to gateways and preserve targets on removal

* test: add unit tests for unassigned target assignment and gateway removal

- getUnassignedTargets: returns targets, empty when no config, empty when field missing
- createGatewayFromWizard: moves selected targets to new gateway, removes from unassigned
- removeGateway: preserves targets as unassigned on removal, no-op for empty gateways
- previewRemoveGateway: shows 'will become unassigned' warning

* style: fix formatting and merge duplicate imports

* docs: add comment explaining unassigned targets preservation

Author: aidandaly24

src/cli/operations/mcp/__tests__/create-mcp.test.ts

@@ -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(() => ({
@@ -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);
+  });
+});

src/cli/operations/mcp/create-mcp.ts

@@ -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.
  */
@@ -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 };

src/cli/operations/remove/__tests__/remove-gateway.test.ts

@@ -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);
+  });
+});

src/cli/operations/remove/remove-gateway.ts

@@ -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
@@ -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] }
+      : {}),
   };
 }
 

src/cli/tui/hooks/useCreateMcp.ts

@@ -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';
@@ -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 };
+}

src/cli/tui/screens/mcp/AddGatewayFlow.tsx

@@ -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';
@@ -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
@@ -157,6 +158,7 @@ export function AddGatewayFlow({
         <AddGatewayScreen
           existingGateways={existingGateways}
           availableAgents={availableAgents}
+          unassignedTargets={unassignedTargets}
           onComplete={handleCreateComplete}
           onExit={onBack}
         />
@@ -183,6 +185,7 @@ export function AddGatewayFlow({
       <AddGatewayScreen
         existingGateways={existingGateways}
         availableAgents={availableAgents}
+        unassignedTargets={unassignedTargets}
         onComplete={handleCreateComplete}
         onExit={() => setFlow({ name: 'mode-select' })}
       />