add rollback support to installation steps

Introduce rollback functionality to the installation process, enabling the reversal of completed steps upon failure. Extend «InstallationAIAssistantStep» with a new «rollback» method and update «InstallationProgressManager» to manage rollback execution. Add rollback handling for agent, connector, and model creation steps, ensuring created resources are cleaned up if subsequent steps fail.

Enhance integration tests to validate rollback behavior and update the «InstallationResult» structure to include «rollbackErrors» for improved error reporting.

Author: guidomodarelli

public/dashboard-assistant/modules/agent/application/ports/__mocks__/agent-repository.ts

@@ -8,6 +8,7 @@ import type { AgentRepository } from '../agent-repository';
 export function createAgentRepositoryMock(): jest.Mocked<AgentRepository> {
   return {
     create: jest.fn(),
+    delete: jest.fn(),
     execute: jest.fn(),
     getActive: jest.fn(),
     register: jest.fn(),

public/dashboard-assistant/modules/agent/application/ports/agent-repository.ts

@@ -3,11 +3,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { CreateRepository } from '../../../common/domain/entities/repository';
+import { CreateRepository, DeleteRepository } from '../../../common/domain/entities/repository';
 import { Agent } from '../../domain/entities/agent';
 import { CreateAgentDto } from '../dtos/create-agent-dto';
 
-export interface AgentRepository extends CreateRepository<Agent, CreateAgentDto> {
+export interface AgentRepository
+  extends CreateRepository<Agent, CreateAgentDto>,
+    DeleteRepository {
   execute(id: string, parameters: any): Promise<any>;
   getActive(): Promise<string | undefined>;
   register(agentId: string): Promise<void>;

public/dashboard-assistant/modules/agent/application/use-cases/delete-agent.ts

@@ -0,0 +1,13 @@
+/*
+ * Copyright Wazuh Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { AgentRepository } from '../ports/agent-repository';
+
+export const deleteAgentUseCase = (agentRepository: AgentRepository) => async (
+  agentId: string
+): Promise<void> => {
+  await agentRepository.delete(agentId);
+};
+

public/dashboard-assistant/modules/connector/application/use-cases/delete-connector.ts

@@ -0,0 +1,13 @@
+/*
+ * Copyright Wazuh Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { ConnectorRepository } from '../ports/connector-repository';
+
+export const deleteConnectorUseCase = (
+  connectorRepository: ConnectorRepository
+) => async (connectorId: string): Promise<void> => {
+  await connectorRepository.delete(connectorId);
+};
+

public/dashboard-assistant/modules/installation-manager/domain/entities/installation-ai-assistant-step.ts

@@ -22,4 +22,9 @@ export abstract class InstallationAIAssistantStep {
   ): Promise<void>;
   abstract getSuccessMessage(): string;
   abstract getFailureMessage(): string;
+  abstract rollback(
+    request: InstallAIDashboardAssistantDto,
+    context: InstallationContext,
+    error: Error
+  ): Promise<void>;
 }

public/dashboard-assistant/modules/installation-manager/domain/entities/installation-context.ts

@@ -21,6 +21,10 @@ export class InstallationContext {
     return this.context.has(key);
   }
 
+  public delete(key: string): void {
+    this.context.delete(key);
+  }
+
   public clear(): void {
     this.context.clear();
   }

public/dashboard-assistant/modules/installation-manager/domain/entities/installation-progress-manager.test.ts

@@ -5,7 +5,9 @@
 
 import { InstallationProgressManager } from './installation-progress-manager';
 import { InstallationAIAssistantStep } from './installation-ai-assistant-step';
-import { ExecutionState, StepResultState } from '../enums';
+import { ExecutionState } from '../enums';
+import type { InstallAIDashboardAssistantDto } from '../types/install-ai-dashboard-assistant-dto';
+import { InstallationContext } from './installation-context';
 
 class TestStep extends InstallationAIAssistantStep {
   constructor(
@@ -24,9 +26,24 @@ class TestStep extends InstallationAIAssistantStep {
   getFailureMessage(): string {
     return this.failureMsg;
   }
+  async rollback(
+    _request: InstallAIDashboardAssistantDto,
+    _context: InstallationContext,
+    _error: Error
+  ): Promise<void> {
+    // no-op for tests
+  }
 }
 
 describe('InstallationProgressManager', () => {
+  const request: InstallAIDashboardAssistantDto = {
+    selected_provider: 'test-provider',
+    model_id: 'test-model',
+    api_url: 'https://example.org',
+    api_key: 'secret',
+  };
+  const createContext = () => new InstallationContext();
+
   it('initializes progress with provided steps in PENDING state', () => {
     const steps = [new TestStep('Step 1'), new TestStep('Step 2')];
     const mgr = new InstallationProgressManager(steps);
@@ -42,7 +59,7 @@ describe('InstallationProgressManager', () => {
     const onProgressChange = jest.fn();
     const mgr = new InstallationProgressManager(steps, onProgressChange);
 
-    await mgr.runStep(steps[0], async () => {
+    await mgr.runStep(steps[0], request, createContext(), async () => {
       /* success */
     });
 
@@ -58,7 +75,7 @@ describe('InstallationProgressManager', () => {
     const mgr = new InstallationProgressManager(steps);
 
     await expect(
-      mgr.runStep(steps[0], async () => {
+      mgr.runStep(steps[0], request, createContext(), async () => {
         throw new Error('boom');
       })
     ).rejects.toThrow('boom');
@@ -76,9 +93,14 @@ describe('InstallationProgressManager', () => {
     const mgr = new InstallationProgressManager(steps);
 
     let resolveExec!: () => void;
-    const running = mgr.runStep(steps[0], () => new Promise<void>((res) => (resolveExec = res)));
+    const running = mgr.runStep(
+      steps[0],
+      request,
+      createContext(),
+      () => new Promise<void>((res) => (resolveExec = res))
+    );
 
-    await expect(mgr.runStep(steps[0], async () => {})).rejects.toThrow();
+    await expect(mgr.runStep(steps[0], request, createContext(), async () => {})).rejects.toThrow();
 
     resolveExec();
     await running;
@@ -88,13 +110,13 @@ describe('InstallationProgressManager', () => {
     const steps = [new TestStep('S1'), new TestStep('S2')];
     const mgr = new InstallationProgressManager(steps);
 
-    await mgr.runStep(steps[0], async () => {});
-    await mgr.runStep(steps[1], async () => {});
+    await mgr.runStep(steps[0], request, createContext(), async () => {});
+    await mgr.runStep(steps[1], request, createContext(), async () => {});
 
     const p = mgr.getProgress();
 
     expect(p.isFinished()).toBe(true);
-    await expect(mgr.runStep(steps[0], async () => {})).rejects.toThrow();
+    await expect(mgr.runStep(steps[0], request, createContext(), async () => {})).rejects.toThrow();
   });
 
   it('reset returns all steps to PENDING and clears result/message/error', async () => {
@@ -103,9 +125,9 @@ describe('InstallationProgressManager', () => {
     const mgr = new InstallationProgressManager(steps, onProgressChange);
 
     // Make one success and one failure
-    await mgr.runStep(steps[0], async () => {});
+    await mgr.runStep(steps[0], request, createContext(), async () => {});
     await expect(
-      mgr.runStep(steps[1], async () => {
+      mgr.runStep(steps[1], request, createContext(), async () => {
         throw new Error('x');
       })
     ).rejects.toThrow('x');

public/dashboard-assistant/modules/installation-manager/domain/entities/installation-progress-manager.ts

@@ -6,11 +6,20 @@
 import { ExecutionState, StepResultState } from '../enums';
 import { InstallationAIAssistantStep } from './installation-ai-assistant-step';
 import { InstallationProgress } from './installation-progress';
+import type { InstallAIDashboardAssistantDto } from '../types/install-ai-dashboard-assistant-dto';
+import type { InstallationContext } from './installation-context';
+
+type RollbackError = {
+  step: string;
+  message: string;
+};
 
 export class InstallationProgressManager {
   private readonly progress: InstallationProgress;
   // Prevent concurrent executions
   private inProgress = false;
+  private completedSteps: InstallationAIAssistantStep[] = [];
+  private rollbackErrors: RollbackError[] = [];
 
   constructor(
     steps: InstallationAIAssistantStep[],
@@ -39,6 +48,8 @@ export class InstallationProgressManager {
 
   public async runStep(
     step: InstallationAIAssistantStep,
+    request: InstallAIDashboardAssistantDto,
+    context: InstallationContext,
     executor: () => Promise<void>
   ): Promise<void> {
     if (this.inProgress) {
@@ -51,18 +62,25 @@ export class InstallationProgressManager {
 
     this.inProgress = true;
     this.progress.startStep(i);
+    this.rollbackErrors = [];
     try {
       await executor();
+      this.completedSteps.push(step);
       this.succeedStep(i, step);
     } catch (err) {
       const error = err instanceof Error ? err : new Error(String(err));
+      await this.rollbackSteps(step, request, context, error);
       this.failStep(i, step, error);
       throw error;
     } finally {
       this.inProgress = false;
     }
   }
 
+  public getRollbackErrors(): RollbackError[] | undefined {
+    return this.rollbackErrors.length > 0 ? [...this.rollbackErrors] : undefined;
+  }
+
   private succeedStep(stepIndex: number, step: InstallationAIAssistantStep): void {
     this.progress.completeStep(stepIndex, StepResultState.SUCCESS, step.getSuccessMessage());
   }
@@ -76,4 +94,34 @@ export class InstallationProgressManager {
       this.onProgressChange(this.getProgress());
     }
   }
+
+  private async rollbackSteps(
+    failedStep: InstallationAIAssistantStep,
+    request: InstallAIDashboardAssistantDto,
+    context: InstallationContext,
+    failure: Error
+  ): Promise<void> {
+    const stepsToRollback = [...this.completedSteps].reverse();
+
+    await this.invokeRollback(failedStep, request, context, failure);
+    for (const step of stepsToRollback) {
+      await this.invokeRollback(step, request, context, failure);
+    }
+
+    this.completedSteps = [];
+  }
+
+  private async invokeRollback(
+    step: InstallationAIAssistantStep,
+    request: InstallAIDashboardAssistantDto,
+    context: InstallationContext,
+    failure: Error
+  ): Promise<void> {
+    try {
+      await step.rollback(request, context, failure);
+    } catch (error) {
+      const normalizedError = error instanceof Error ? error : new Error(String(error));
+      this.rollbackErrors.push({ step: step.getName(), message: normalizedError.message });
+    }
+  }
 }

public/dashboard-assistant/modules/installation-manager/domain/types/installation-result.ts

@@ -18,4 +18,8 @@ export interface InstallationResult {
     [key: string]: any;
   };
   errors?: InstallationError[];
+  rollbackErrors?: Array<{
+    step: string;
+    message: string;
+  }>;
 }

public/dashboard-assistant/modules/installation-manager/infrastructure/installation-manager.int.test.ts

@@ -38,6 +38,9 @@ describe('InstallationManager (integration)', () => {
       validateModelConnection: jest.fn().mockResolvedValue(true),
       createAgent: jest.fn().mockResolvedValue({ id: 'agent-1' }),
       useAgent: jest.fn().mockResolvedValue(undefined),
+      deleteConnector: jest.fn(),
+      deleteModel: jest.fn(),
+      deleteAgent: jest.fn(),
     };
 
     const progressUpdates: unknown[] = [];
@@ -70,14 +73,97 @@ describe('InstallationManager (integration)', () => {
       validateModelConnection: jest.fn(),
       createAgent: jest.fn(),
       useAgent: jest.fn(),
+      deleteConnector: jest.fn().mockResolvedValue(undefined),
+      deleteModel: jest.fn(),
+      deleteAgent: jest.fn(),
     };
 
     const manager = new InstallationManager();
     const result = await manager.execute(request);
     expect(result.success).toBe(false);
     expect(result.errors).toBeDefined();
     expect(result.errors![0]).toMatchObject({ step: 'Create Model' });
-    // Data should contain only the connectorId
-    expect(result.data).toEqual({ connectorId: 'conn-1' });
+    // Data should not contain identifiers for successfully rolled-back steps
+    expect(result.data).toEqual({});
+    expect(result.rollbackErrors).toBeUndefined();
+    expect(
+      ((global as unknown) as {
+        __mockUseCases: import('../../../services/__mocks__').MockUseCases;
+      }).__mockUseCases!.deleteConnector
+    ).toHaveBeenCalledWith('conn-1');
+    expect(
+      ((global as unknown) as {
+        __mockUseCases: import('../../../services/__mocks__').MockUseCases;
+      }).__mockUseCases!.deleteModel
+    ).not.toHaveBeenCalled();
+    expect(
+      ((global as unknown) as {
+        __mockUseCases: import('../../../services/__mocks__').MockUseCases;
+      }).__mockUseCases!.deleteAgent
+    ).not.toHaveBeenCalled();
+  });
+
+  it('rolls back created resources when a later step fails', async () => {
+    ((global as unknown) as {
+      __mockUseCases: import('../../../services/__mocks__').MockUseCases;
+    }).__mockUseCases = {
+      persistMlCommonsSettings: jest.fn().mockResolvedValue(undefined),
+      createConnector: jest.fn().mockResolvedValue({ id: 'conn-1' }),
+      createModel: jest.fn().mockResolvedValue({ id: 'model-1' }),
+      validateModelConnection: jest.fn().mockResolvedValue(true),
+      createAgent: jest.fn().mockResolvedValue({ id: 'agent-1' }),
+      useAgent: jest.fn().mockRejectedValue(new Error('register failed')),
+      deleteConnector: jest.fn().mockResolvedValue(undefined),
+      deleteModel: jest.fn().mockResolvedValue(undefined),
+      deleteAgent: jest.fn().mockResolvedValue(undefined),
+    };
+
+    const manager = new InstallationManager();
+    const result = await manager.execute(request);
+
+    expect(result.success).toBe(false);
+    expect(result.errors![0]).toMatchObject({ step: 'Register Agent' });
+    expect(result.data).toEqual({});
+    expect(result.rollbackErrors).toBeUndefined();
+
+    const mockUseCases = ((global as unknown) as {
+      __mockUseCases: import('../../../services/__mocks__').MockUseCases;
+    }).__mockUseCases!;
+
+    expect(mockUseCases.deleteAgent).toHaveBeenCalledWith('agent-1');
+    expect(mockUseCases.deleteModel).toHaveBeenCalledWith('model-1');
+    expect(mockUseCases.deleteConnector).toHaveBeenCalledWith('conn-1');
+  });
+
+  it('rolls back created resources when model connection validation fails', async () => {
+    ((global as unknown) as {
+      __mockUseCases: import('../../../services/__mocks__').MockUseCases;
+    }).__mockUseCases = {
+      persistMlCommonsSettings: jest.fn().mockResolvedValue(undefined),
+      createConnector: jest.fn().mockResolvedValue({ id: 'conn-1' }),
+      createModel: jest.fn().mockResolvedValue({ id: 'model-1' }),
+      validateModelConnection: jest.fn().mockRejectedValue(new Error('validation failed')),
+      createAgent: jest.fn(),
+      useAgent: jest.fn(),
+      deleteConnector: jest.fn().mockResolvedValue(undefined),
+      deleteModel: jest.fn().mockResolvedValue(undefined),
+      deleteAgent: jest.fn().mockResolvedValue(undefined),
+    };
+
+    const manager = new InstallationManager();
+    const result = await manager.execute(request);
+
+    expect(result.success).toBe(false);
+    expect(result.errors![0]).toMatchObject({ step: 'Test Model Connection' });
+    expect(result.data).toEqual({});
+    expect(result.rollbackErrors).toBeUndefined();
+
+    const mockUseCases = ((global as unknown) as {
+      __mockUseCases: import('../../../services/__mocks__').MockUseCases;
+    }).__mockUseCases!;
+
+    expect(mockUseCases.deleteModel).toHaveBeenCalledWith('model-1');
+    expect(mockUseCases.deleteConnector).toHaveBeenCalledWith('conn-1');
+    expect(mockUseCases.deleteAgent).not.toHaveBeenCalled();
   });
 });