Skip to content
Draft
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
14 changes: 12 additions & 2 deletions integ-tests/add-remove-gateway.test.ts
Original file line number Diff line number Diff line change
@@ -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'));
}
Expand All @@ -19,6 +22,7 @@ describe('integration: add and remove gateway with external MCP server', () => {

afterAll(async () => {
await project.cleanup();
telemetry.destroy();
});

describe('gateway lifecycle', () => {
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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' });
});
});
});
Expand Down
45 changes: 39 additions & 6 deletions integ-tests/add-remove-resources.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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' });
});

Expand Down Expand Up @@ -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',
Expand All @@ -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);
Expand All @@ -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' });
});
});

Expand All @@ -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',
Expand All @@ -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);
Expand All @@ -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' });
});
});

Expand All @@ -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' });
});
});
});
10 changes: 7 additions & 3 deletions src/cli/commands/remove/command.tsx
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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 => {
Expand Down
8 changes: 7 additions & 1 deletion src/cli/primitives/BasePrimitive.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/cli/primitives/GatewayPrimitive.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/cli/primitives/GatewayTargetPrimitive.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/cli/primitives/PolicyEnginePrimitive.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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({
Expand Down
4 changes: 2 additions & 2 deletions src/cli/primitives/PolicyPrimitive.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(
Expand Down
4 changes: 2 additions & 2 deletions src/cli/primitives/RuntimeEndpointPrimitive.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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,
Expand Down
25 changes: 17 additions & 8 deletions src/cli/telemetry/cli-command-run.ts
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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();
Expand All @@ -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!;
Expand Down
18 changes: 18 additions & 0 deletions src/cli/telemetry/schemas/command-run.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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('.');
Expand Down
Loading
Loading