feat: wire telemetry into all remove.* commands

Hweinstock · Hweinstock · commit 304bda5c8469 · 2026-04-30T22:19:01.000Z
diff --git a/integ-tests/add-remove-resources.test.ts b/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,13 @@ 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' });
     });
   });
 });
diff --git a/src/cli/commands/remove/command.tsx b/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 => {
diff --git a/src/cli/primitives/BasePrimitive.ts b/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,
diff --git a/src/cli/primitives/GatewayPrimitive.ts b/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,
diff --git a/src/cli/primitives/GatewayTargetPrimitive.ts b/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,
diff --git a/src/cli/primitives/PolicyEnginePrimitive.ts b/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({
diff --git a/src/cli/primitives/PolicyPrimitive.ts b/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(
diff --git a/src/cli/primitives/RuntimeEndpointPrimitive.ts b/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,
diff --git a/src/cli/telemetry/cli-command-run.ts b/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!;
diff --git a/src/cli/telemetry/schemas/command-run.ts b/src/cli/telemetry/schemas/command-run.ts
@@ -195,6 +195,7 @@ export const COMMAND_SCHEMAS = {
   validate: NoAttrs,
   'help.modes': NoAttrs,
   help: NoAttrs,
+  'remove.all': NoAttrs,
   'remove.agent': NoAttrs,
   'remove.memory': NoAttrs,
   'remove.credential': NoAttrs,
@@ -204,6 +205,7 @@ export const COMMAND_SCHEMAS = {
   'remove.gateway-target': NoAttrs,
   'remove.policy-engine': NoAttrs,
   'remove.policy': NoAttrs,
+  'remove.runtime-endpoint': NoAttrs,
   'telemetry.disable': NoAttrs,
   'telemetry.enable': NoAttrs,
   'telemetry.status': NoAttrs,
@@ -216,6 +218,22 @@ export const COMMAND_SCHEMAS = {
 export type Command = keyof typeof COMMAND_SCHEMAS;
 export type CommandAttrs<C extends Command> = z.infer<(typeof COMMAND_SCHEMAS)[C]>;
 
+/** Extract the command group prefix from a dotted command key (e.g. 'add' from 'add.agent'). */
+type CommandGroup = {
+  [C in Command]: C extends `${infer G}.${string}` ? G : C;
+}[Command];
+
+/**
+ * Type-safe lookup of a subcommand under a command group.
+ * Produces a compile-time error if `${G}.${S}` is not a registered command.
+ *
+ * @example
+ * SubCommand<'remove', 'agent'>  // → 'remove.agent'
+ * SubCommand<'add', 'memory'>    // → 'add.memory'
+ * SubCommand<'remove', 'bogus'>  // → never (compile error at call site)
+ */
+export type SubCommand<G extends CommandGroup, S extends string> = Extract<Command, `${G}.${S}`>;
+
 /** Derive command_group from command key (e.g. 'add.agent' → 'add') */
 export function deriveCommandGroup(command: Command): string {
   const dot = command.indexOf('.');
diff --git a/src/cli/tui/hooks/useCreateEvaluator.ts b/src/cli/tui/hooks/useCreateEvaluator.ts
@@ -1,6 +1,6 @@
 import type { EvaluatorConfig } from '../../../schema';
 import { evaluatorPrimitive } from '../../primitives/registry';
-import { withAddTelemetry } from '../../telemetry/cli-command-run.js';
+import { withTuiTelemetry } from '../../telemetry/cli-command-run.js';
 import { Level, standardize } from '../../telemetry/schemas/common-shapes.js';
 import { useCallback, useEffect, useState } from 'react';
 
@@ -18,7 +18,7 @@ export function useCreateEvaluator() {
   const create = useCallback(async (config: CreateEvaluatorConfig) => {
     setStatus({ state: 'loading' });
     try {
-      const addResult = await withAddTelemetry(
+      const addResult = await withTuiTelemetry(
         'add.evaluator',
         {
           evaluator_type: config.config.codeBased ? 'code-based' : 'llm-as-a-judge',
diff --git a/src/cli/tui/hooks/useCreateMcp.ts b/src/cli/tui/hooks/useCreateMcp.ts
@@ -4,7 +4,7 @@ import {
   gatewayTargetPrimitive,
   policyEnginePrimitive,
 } from '../../primitives/registry';
-import { withAddTelemetry } from '../../telemetry/cli-command-run.js';
+import { withTuiTelemetry } from '../../telemetry/cli-command-run.js';
 import { AuthorizerType, PolicyEngineMode, standardize } from '../../telemetry/schemas/common-shapes.js';
 import type { AddGatewayConfig } from '../screens/mcp/types';
 import { useCallback, useEffect, useState } from 'react';
@@ -25,7 +25,7 @@ export function useCreateGateway() {
   const createGateway = useCallback(async (config: AddGatewayConfig) => {
     setStatus({ state: 'loading' });
     try {
-      const addResult = await withAddTelemetry(
+      const addResult = await withTuiTelemetry(
         'add.gateway',
         {
           authorizer_type: standardize(AuthorizerType, config.authorizerType ?? 'NONE'),
diff --git a/src/cli/tui/hooks/useCreateMemory.ts b/src/cli/tui/hooks/useCreateMemory.ts
@@ -2,7 +2,7 @@ import { ConfigIO } from '../../../lib';
 import type { Memory } from '../../../schema';
 import { getAvailableAgents } from '../../operations/attach';
 import { memoryPrimitive } from '../../primitives/registry';
-import { withAddTelemetry } from '../../telemetry/cli-command-run.js';
+import { withTuiTelemetry } from '../../telemetry/cli-command-run.js';
 import { useCallback, useEffect, useState } from 'react';
 
 interface CreateMemoryConfig {
@@ -26,7 +26,7 @@ export function useCreateMemory() {
     try {
       const strategiesStr = config.strategies.map(s => s.type).join(',');
       const strategyList = strategiesStr ? strategiesStr.split(',').map(s => s.trim().toUpperCase()) : [];
-      const addResult = await withAddTelemetry(
+      const addResult = await withTuiTelemetry(
         'add.memory',
         {
           strategy_count: strategyList.length,
diff --git a/src/cli/tui/hooks/useCreateOnlineEval.ts b/src/cli/tui/hooks/useCreateOnlineEval.ts
@@ -1,5 +1,5 @@
 import { onlineEvalConfigPrimitive } from '../../primitives/registry';
-import { withAddTelemetry } from '../../telemetry/cli-command-run.js';
+import { withTuiTelemetry } from '../../telemetry/cli-command-run.js';
 import { useCallback, useEffect, useState } from 'react';
 
 interface CreateOnlineEvalConfig {
@@ -18,7 +18,7 @@ export function useCreateOnlineEval() {
   const create = useCallback(async (config: CreateOnlineEvalConfig) => {
     setStatus({ state: 'loading' });
     try {
-      const addResult = await withAddTelemetry(
+      const addResult = await withTuiTelemetry(
         'add.online-eval',
         {
           evaluator_count: config.evaluators.length,
diff --git a/src/cli/tui/hooks/useRemove.ts b/src/cli/tui/hooks/useRemove.ts
diff --git a/src/cli/tui/screens/agent/useAddAgent.ts b/src/cli/tui/screens/agent/useAddAgent.ts
diff --git a/src/cli/tui/screens/identity/useCreateIdentity.ts b/src/cli/tui/screens/identity/useCreateIdentity.ts
diff --git a/src/cli/tui/screens/policy/AddPolicyFlow.tsx b/src/cli/tui/screens/policy/AddPolicyFlow.tsx
diff --git a/src/cli/tui/screens/remove/useRemoveFlow.ts b/src/cli/tui/screens/remove/useRemoveFlow.ts
diff --git a/src/cli/tui/screens/runtime-endpoint/AddRuntimeEndpointFlow.tsx b/src/cli/tui/screens/runtime-endpoint/AddRuntimeEndpointFlow.tsx
