fix: withAddTelemetry defensive init, PolicyEnginePrimitive keeps telemetry in CLI handler

Author: Hweinstock

src/cli/primitives/PolicyEnginePrimitive.ts

@@ -3,7 +3,10 @@ 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, withAddTelemetry } from '../telemetry/cli-command-run.js';
+import { cliCommandRun } from '../telemetry/cli-command-run.js';
+import { TelemetryClientAccessor } from '../telemetry/client-accessor.js';
+import type { CommandAttrs } from '../telemetry/schemas/command-run.js';
+import { AttachMode, standardize } from '../telemetry/schemas/common-shapes.js';
 import { requireTTY } from '../tui/guards/tty';
 import { BasePrimitive } from './BasePrimitive';
 import { SOURCE_CODE_NOTE } from './constants';
@@ -22,27 +25,25 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
   readonly primitiveSchema = PolicyEngineSchema;
 
   async add(options: AddPolicyEngineOptions): Promise<AddResult<{ engineName: string }>> {
-    return withAddTelemetry('add.policy-engine', { attach_gateway_count: 0, attach_mode: 'log_only' }, async () => {
-      try {
-        const project = await this.readProjectSpec();
+    try {
+      const project = await this.readProjectSpec();
 
-        this.checkDuplicate(project.policyEngines, options.name);
+      this.checkDuplicate(project.policyEngines, options.name);
 
-        const engine: PolicyEngine = {
-          name: options.name,
-          ...(options.description && { description: options.description }),
-          ...(options.encryptionKeyArn && { encryptionKeyArn: options.encryptionKeyArn }),
-          policies: [],
-        };
+      const engine: PolicyEngine = {
+        name: options.name,
+        ...(options.description && { description: options.description }),
+        ...(options.encryptionKeyArn && { encryptionKeyArn: options.encryptionKeyArn }),
+        policies: [],
+      };
 
-        project.policyEngines.push(engine);
-        await this.writeProjectSpec(project);
+      project.policyEngines.push(engine);
+      await this.writeProjectSpec(project);
 
-        return { success: true, engineName: engine.name };
-      } catch (err) {
-        return { success: false, error: getErrorMessage(err) };
-      }
-    });
+      return { success: true, engineName: engine.name };
+    } catch (err) {
+      return { success: false, error: getErrorMessage(err) };
+    }
   }
 
   async remove(engineName: string): Promise<RemovalResult> {
@@ -235,30 +236,56 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
                 throw new Error('--name is required');
               }
 
-              const result = await this.add({
-                name: cliOptions.name,
-                description: cliOptions.description,
-                encryptionKeyArn: cliOptions.encryptionKeyArn,
-              });
-
-              // Attach to gateways if requested
-              if (result.success && cliOptions.attachToGateways) {
-                const mode = PolicyEngineModeSchema.parse(cliOptions.attachMode ?? 'LOG_ONLY');
-                const gateways = cliOptions.attachToGateways
-                  .split(',')
-                  .map(s => s.trim())
-                  .filter(Boolean);
-                await this.attachToGateways(cliOptions.name, gateways, mode);
+              const gatewayNames = cliOptions.attachToGateways
+                ? cliOptions.attachToGateways
+                    .split(',')
+                    .map(s => s.trim())
+                    .filter(Boolean)
+                : [];
+              const attachMode = cliOptions.attachMode ?? 'LOG_ONLY';
+              const telemetryAttrs: CommandAttrs<'add.policy-engine'> = {
+                attach_gateway_count: gatewayNames.length,
+                attach_mode: standardize(AttachMode, attachMode),
+              };
+
+              let client;
+              try {
+                client = await TelemetryClientAccessor.get();
+              } catch {
+                /* telemetry init failed */
               }
 
-              if (!result.success) {
-                throw new Error(result.error);
-              }
+              const name = cliOptions.name;
+
+              const run = async () => {
+                const result = await this.add({
+                  name,
+                  description: cliOptions.description,
+                  encryptionKeyArn: cliOptions.encryptionKeyArn,
+                });
+
+                if (result.success && gatewayNames.length > 0) {
+                  const mode = PolicyEngineModeSchema.parse(attachMode);
+                  await this.attachToGateways(name, gatewayNames, mode);
+                }
+
+                if (!result.success) {
+                  throw new Error(result.error);
+                }
+
+                if (cliOptions.json) {
+                  console.log(JSON.stringify(result));
+                } else {
+                  console.log(`Added policy engine '${result.engineName}'`);
+                }
+
+                return telemetryAttrs;
+              };
 
-              if (cliOptions.json) {
-                console.log(JSON.stringify(result));
+              if (client) {
+                await client.withCommandRun('add.policy-engine', run);
               } else {
-                console.log(`Added policy engine '${result.engineName}'`);
+                await run();
               }
             });
           } else {

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

@@ -39,15 +39,15 @@ export async function withAddTelemetry<C extends Command, T extends Record<strin
     return fn();
   }
 
-  let result!: AddResult<T>;
+  let result: AddResult<T> | undefined;
   try {
     await client.withCommandRun(command, async () => {
       result = await fn();
       if (!result.success) throw new Error(result.error);
       return attrs;
     });
   } catch {
-    // withCommandRun already recorded the failure; swallow the re-throw
+    if (!result) return fn();
   }
-  return result;
+  return result!;
 }