feat: wire telemetry into all remove.* commands

Author: Hweinstock

integ-tests/add-remove-gateway.test.ts

@@ -1,9 +1,12 @@
 import { createTestProject, runCLI } from '../src/test-utils/index.js';
 import type { TestProject } from '../src/test-utils/index.js';
+import { createTelemetryHelper } from '../src/test-utils/telemetry-helper.js';
 import { mkdir, readFile, writeFile } from 'node:fs/promises';
 import { join } from 'node:path';
 import { afterAll, beforeAll, describe, expect, it } from 'vitest';
 
+const telemetry = createTelemetryHelper();
+
 async function readProjectConfig(projectPath: string) {
   return JSON.parse(await readFile(join(projectPath, 'agentcore/agentcore.json'), 'utf-8'));
 }
@@ -19,6 +22,7 @@ describe('integration: add and remove gateway with external MCP server', () => {
 
   afterAll(async () => {
     await project.cleanup();
+    telemetry.destroy();
   });
 
   describe('gateway lifecycle', () => {
@@ -64,7 +68,9 @@ describe('integration: add and remove gateway with external MCP server', () => {
     });
 
     it('removes the gateway target', async () => {
-      const result = await runCLI(['remove', 'gateway-target', '--name', targetName, '--json'], project.projectPath);
+      const result = await runCLI(['remove', 'gateway-target', '--name', targetName, '--json'], project.projectPath, {
+        env: telemetry.env,
+      });
 
       expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
       const json = JSON.parse(result.stdout);
@@ -75,10 +81,13 @@ describe('integration: add and remove gateway with external MCP server', () => {
       const targets = gateway?.targets ?? [];
       const found = targets.find((t: { name: string }) => t.name === targetName);
       expect(found, `Target "${targetName}" should be removed`).toBeFalsy();
+      telemetry.assertMetricEmitted({ command: 'remove.gateway-target', exit_reason: 'success' });
     });
 
     it('removes the gateway', async () => {
-      const result = await runCLI(['remove', 'gateway', '--name', gatewayName, '--json'], project.projectPath);
+      const result = await runCLI(['remove', 'gateway', '--name', gatewayName, '--json'], project.projectPath, {
+        env: telemetry.env,
+      });
 
       expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
       const json = JSON.parse(result.stdout);
@@ -88,6 +97,7 @@ describe('integration: add and remove gateway with external MCP server', () => {
       const gateways = mcpSpec.agentCoreGateways ?? [];
       const found = gateways.find((g: { name: string }) => g.name === gatewayName);
       expect(found, `Gateway "${gatewayName}" should be removed`).toBeFalsy();
+      telemetry.assertMetricEmitted({ command: 'remove.gateway', exit_reason: 'success' });
     });
   });
 });

integ-tests/add-remove-resources.test.ts

@@ -41,7 +41,6 @@ describe('integration: add and remove resources', () => {
       const found = memories!.some((m: Record<string, unknown>) => m.name === memoryName);
       expect(found, `Memory "${memoryName}" should be in config`).toBe(true);
 
-      // Verify telemetry
       telemetry.assertMetricEmitted({ command: 'add.memory', exit_reason: 'success' });
     });
 
@@ -71,7 +70,6 @@ describe('integration: add and remove resources', () => {
       expect(episodic!.reflectionNamespaces, 'Should have reflectionNamespaces').toBeDefined();
       expect(episodic!.reflectionNamespaces!.length).toBeGreaterThan(0);
 
-      // Verify telemetry
       telemetry.assertMetricEmitted({
         command: 'add.memory',
         exit_reason: 'success',
@@ -84,7 +82,9 @@ describe('integration: add and remove resources', () => {
     });
 
     it('removes the memory resource', async () => {
-      const result = await runCLI(['remove', 'memory', '--name', memoryName, '--json'], project.projectPath);
+      const result = await runCLI(['remove', 'memory', '--name', memoryName, '--json'], project.projectPath, {
+        env: telemetry.env,
+      });
 
       expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
       const json = JSON.parse(result.stdout);
@@ -95,6 +95,8 @@ describe('integration: add and remove resources', () => {
       const memories = (config.memories as Record<string, unknown>[] | undefined) ?? [];
       const found = memories.some((m: Record<string, unknown>) => m.name === memoryName);
       expect(found, `Memory "${memoryName}" should be removed from config`).toBe(false);
+
+      telemetry.assertMetricEmitted({ command: 'remove.memory', exit_reason: 'success' });
     });
   });
 
@@ -119,7 +121,6 @@ describe('integration: add and remove resources', () => {
       const found = credentials!.some((c: Record<string, unknown>) => c.name === credentialName);
       expect(found, `Credential "${credentialName}" should be in config`).toBe(true);
 
-      // Verify telemetry
       telemetry.assertMetricEmitted({
         command: 'add.credential',
         exit_reason: 'success',
@@ -128,7 +129,9 @@ describe('integration: add and remove resources', () => {
     });
 
     it('removes the credential resource', async () => {
-      const result = await runCLI(['remove', 'credential', '--name', credentialName, '--json'], project.projectPath);
+      const result = await runCLI(['remove', 'credential', '--name', credentialName, '--json'], project.projectPath, {
+        env: telemetry.env,
+      });
 
       expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
       const json = JSON.parse(result.stdout);
@@ -139,6 +142,8 @@ describe('integration: add and remove resources', () => {
       const credentials = (config.credentials as Record<string, unknown>[] | undefined) ?? [];
       const found = credentials.some((c: Record<string, unknown>) => c.name === credentialName);
       expect(found, `Credential "${credentialName}" should be removed from config`).toBe(false);
+
+      telemetry.assertMetricEmitted({ command: 'remove.credential', exit_reason: 'success' });
     });
   });
 
@@ -162,9 +167,37 @@ describe('integration: add and remove resources', () => {
     });
 
     it('removes the policy engine resource', async () => {
-      const result = await runCLI(['remove', 'policy-engine', '--name', engineName, '--json'], project.projectPath);
+      const result = await runCLI(['remove', 'policy-engine', '--name', engineName, '--json'], project.projectPath, {
+        env: telemetry.env,
+      });
 
       expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0);
+
+      telemetry.assertMetricEmitted({ command: 'remove.policy-engine', exit_reason: 'success' });
+    });
+  });
+
+  describe('remove failure telemetry', () => {
+    it('emits failure telemetry for non-existent resource', async () => {
+      const result = await runCLI(['remove', 'memory', '--name', 'DoesNotExist', '--json'], project.projectPath, {
+        env: telemetry.env,
+      });
+
+      expect(result.exitCode).toBe(1);
+      telemetry.assertMetricEmitted({ command: 'remove.memory', exit_reason: 'failure' });
+    });
+  });
+
+  describe('remove all', () => {
+    it('resets all schemas and emits telemetry', async () => {
+      const result = await runCLI(['remove', 'all', '--yes', '--json'], project.projectPath, {
+        env: telemetry.env,
+      });
+
+      expect(result.exitCode).toBe(0);
+      const json = JSON.parse(result.stdout);
+      expect(json.success).toBe(true);
+      telemetry.assertMetricEmitted({ command: 'remove.all', exit_reason: 'success' });
     });
   });
 });

src/cli/commands/remove/command.tsx

@@ -1,5 +1,6 @@
 import { ConfigIO } from '../../../lib';
 import { getErrorMessage } from '../../errors';
+import { cliCommandRun } from '../../telemetry/cli-command-run.js';
 import { COMMAND_DESCRIPTIONS } from '../../tui/copy';
 import { requireProject, requireTTY } from '../../tui/guards';
 import { RemoveAllScreen, RemoveFlow } from '../../tui/screens/remove';
@@ -51,9 +52,12 @@ async function handleRemoveAll(_options: RemoveAllOptions): Promise<RemoveResult
 
 async function handleRemoveAllCLI(options: RemoveAllOptions): Promise<void> {
   validateRemoveAllOptions(options);
-  const result = await handleRemoveAll(options);
-  console.log(JSON.stringify(result));
-  process.exit(result.success ? 0 : 1);
+  await cliCommandRun('remove.all', !!options.json, async () => {
+    const result = await handleRemoveAll(options);
+    if (!result.success) throw new Error(result.error);
+    console.log(JSON.stringify(result));
+    return {};
+  });
 }
 
 export const registerRemove = (program: Command): Command => {

src/cli/primitives/BasePrimitive.ts

@@ -2,6 +2,8 @@ import { ConfigIO, findConfigRoot } from '../../lib';
 import type { AgentCoreProjectSpec } from '../../schema';
 import type { ResourceType } from '../commands/remove/types';
 import { getErrorMessage } from '../errors';
+import { withTuiTelemetry } from '../telemetry/cli-command-run.js';
+import type { SubCommand } from '../telemetry/schemas/command-run.js';
 import { requireTTY } from '../tui/guards/tty';
 import { SOURCE_CODE_NOTE } from './constants';
 import type { AddResult, AddScreenComponent, RemovableResource, RemovalPreview, RemovalResult } from './types';
@@ -120,7 +122,11 @@ export abstract class BasePrimitive<
               process.exit(1);
             }
 
-            const result = await this.remove(cliOptions.name);
+            const result = await withTuiTelemetry<SubCommand<'remove', typeof this.kind>, RemovalResult>(
+              `remove.${this.kind}`,
+              {},
+              () => this.remove(cliOptions.name!)
+            );
             console.log(
               JSON.stringify({
                 success: result.success,

src/cli/primitives/GatewayPrimitive.ts

@@ -12,7 +12,7 @@ import type { AddGatewayOptions as CLIAddGatewayOptions } from '../commands/add/
 import { validateAddGatewayOptions } from '../commands/add/validate';
 import { getErrorMessage } from '../errors';
 import type { RemovalPreview, RemovalResult, SchemaChange } from '../operations/remove/types';
-import { cliCommandRun } from '../telemetry/cli-command-run.js';
+import { cliCommandRun, withTuiTelemetry } from '../telemetry/cli-command-run.js';
 import { AuthorizerType, PolicyEngineMode, standardize } from '../telemetry/schemas/common-shapes.js';
 import { requireTTY } from '../tui/guards/tty';
 import type { AddGatewayConfig } from '../tui/screens/mcp/types';
@@ -262,7 +262,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
               process.exit(1);
             }
 
-            const result = await this.remove(cliOptions.name);
+            const result = await withTuiTelemetry('remove.gateway', {}, () => this.remove(cliOptions.name!));
             console.log(
               JSON.stringify({
                 success: result.success,

src/cli/primitives/GatewayTargetPrimitive.ts

@@ -14,7 +14,7 @@ import { validateAddGatewayTargetOptions } from '../commands/add/validate';
 import { getErrorMessage } from '../errors';
 import type { RemovableGatewayTarget } from '../operations/remove/remove-gateway-target';
 import type { RemovalPreview, RemovalResult, SchemaChange } from '../operations/remove/types';
-import { cliCommandRun } from '../telemetry/cli-command-run.js';
+import { cliCommandRun, withTuiTelemetry } from '../telemetry/cli-command-run.js';
 import {
   GATEWAY_TARGET_TYPE_MAP,
   GatewayTargetHost,
@@ -510,7 +510,7 @@ export class GatewayTargetPrimitive extends BasePrimitive<AddGatewayTargetOption
               process.exit(1);
             }
 
-            const result = await this.remove(cliOptions.name);
+            const result = await withTuiTelemetry('remove.gateway-target', {}, () => this.remove(cliOptions.name!));
             console.log(
               JSON.stringify({
                 success: result.success,

src/cli/primitives/PolicyEnginePrimitive.ts

@@ -3,7 +3,7 @@ import type { AgentCoreProjectSpec, PolicyEngine } from '../../schema';
 import { PolicyEngineModeSchema, PolicyEngineSchema } from '../../schema';
 import { getErrorMessage } from '../errors';
 import type { RemovalPreview, RemovalResult, SchemaChange } from '../operations/remove/types';
-import { cliCommandRun } from '../telemetry/cli-command-run.js';
+import { cliCommandRun, withTuiTelemetry } from '../telemetry/cli-command-run.js';
 import { AttachMode, standardize } from '../telemetry/schemas/common-shapes.js';
 import { requireTTY } from '../tui/guards/tty';
 import { BasePrimitive } from './BasePrimitive';
@@ -316,7 +316,7 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
               process.exit(1);
             }
 
-            const result = await this.remove(cliOptions.name);
+            const result = await withTuiTelemetry('remove.policy-engine', {}, () => this.remove(cliOptions.name!));
             if (cliOptions.json) {
               console.log(
                 JSON.stringify({

src/cli/primitives/PolicyPrimitive.ts

@@ -5,7 +5,7 @@ import { detectRegion } from '../aws';
 import { getPolicyGeneration, startPolicyGeneration } from '../aws/policy-generation';
 import { getErrorMessage } from '../errors';
 import type { RemovalPreview, RemovalResult, SchemaChange } from '../operations/remove/types';
-import { cliCommandRun } from '../telemetry/cli-command-run.js';
+import { cliCommandRun, withTuiTelemetry } from '../telemetry/cli-command-run.js';
 import { ValidationMode, standardize } from '../telemetry/schemas/common-shapes.js';
 import { requireTTY } from '../tui/guards/tty';
 import { BasePrimitive } from './BasePrimitive';
@@ -400,7 +400,7 @@ export class PolicyPrimitive extends BasePrimitive<AddPolicyOptions, RemovablePo
 
             // Build composite key when --engine is provided for unambiguous removal
             const removeKey = cliOptions.engine ? `${cliOptions.engine}/${cliOptions.name}` : cliOptions.name;
-            const result = await this.remove(removeKey);
+            const result = await withTuiTelemetry('remove.policy', {}, () => this.remove(removeKey));
 
             if (cliOptions.json) {
               console.log(

src/cli/primitives/RuntimeEndpointPrimitive.ts

@@ -4,7 +4,7 @@ import { RuntimeEndpointSchema } from '../../schema';
 import type { ResourceType } from '../commands/remove/types';
 import { getErrorMessage } from '../errors';
 import type { RemovalPreview, RemovalResult, SchemaChange } from '../operations/remove/types';
-import { cliCommandRun } from '../telemetry/cli-command-run.js';
+import { cliCommandRun, withTuiTelemetry } from '../telemetry/cli-command-run.js';
 import { BasePrimitive } from './BasePrimitive';
 import { SOURCE_CODE_NOTE } from './constants';
 import type { AddResult, AddScreenComponent, RemovableResource } from './types';
@@ -296,7 +296,7 @@ export class RuntimeEndpointPrimitive extends BasePrimitive<AddRuntimeEndpointOp
               process.exit(1);
             }
 
-            const result = await this.remove(cliOptions.name);
+            const result = await withTuiTelemetry('remove.runtime-endpoint', {}, () => this.remove(cliOptions.name!));
             console.log(
               JSON.stringify({
                 success: result.success,

src/cli/telemetry/cli-command-run.ts

@@ -1,8 +1,10 @@
 import { getErrorMessage } from '../errors';
-import type { AddResult } from '../primitives/types.js';
 import { TelemetryClientAccessor } from './client-accessor.js';
 import type { Command, CommandAttrs } from './schemas/command-run.js';
 
+/** Standard result shape for operations wrapped by {@link withTuiTelemetry}. */
+export type OperationResult = { success: true } | { success: false; error: string };
+
 /**
  * Run a CLI command with telemetry, standardized error output, and process.exit.
  * The callback should throw on failure and return telemetry attrs on success.
@@ -38,22 +40,29 @@ export async function cliCommandRun<C extends Command>(
 }
 
 /**
- * Wrap a primitive .add() call with telemetry — used by TUI paths.
- * CLI paths use {@link cliCommandRun} instead.
+ * Wrap a TUI operation with telemetry. Works for any `{ success: boolean }` result type
+ * (add, remove, etc.). CLI paths use {@link cliCommandRun} instead.
+ *
+ * Attrs are only recorded on success — on failure, `withCommandRun` records the
+ * error classification instead.
+ *
+ * @param command  Telemetry command key (e.g. 'add.memory', 'remove.agent')
+ * @param attrs    Command-specific telemetry attributes (pass `{}` for NoAttrs commands)
+ * @param fn       The operation to wrap
  */
-export async function withAddTelemetry<C extends Command, T extends Record<string, unknown>>(
+export async function withTuiTelemetry<C extends Command, R extends OperationResult>(
   command: C,
   attrs: CommandAttrs<C>,
-  fn: () => Promise<AddResult<T>>
-): Promise<AddResult<T>> {
+  fn: () => Promise<R>
+): Promise<R> {
   let client;
   try {
     client = await TelemetryClientAccessor.get();
   } catch {
     return fn();
   }
 
-  let result: AddResult<T> | undefined;
+  let result: R | undefined;
   try {
     await client.withCommandRun(command, async () => {
       result = await fn();
@@ -64,7 +73,7 @@ export async function withAddTelemetry<C extends Command, T extends Record<strin
     // withCommandRun re-throws after recording failure telemetry.
     // result is set if fn() ran; if not, fn() itself threw.
     if (!result) {
-      return { success: false, error: getErrorMessage(err) };
+      return { success: false, error: getErrorMessage(err) } as R;
     }
   }
   return result!;