From 4c9dec519b752f0c322f43ca6d2fbb28964469f8 Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 25 Jun 2026 06:23:54 +0000 Subject: [PATCH] feat(CLI telemetry): shared finalize-before-process.exit ordering (1439) + command_run instrumentation for config (1446), run-eval/pause/resume online-eval+online-insights/traces (1447), and the fetch-access TUI path (1448). fetch-access CLI and run.job/ab-test were already instrumented at HEAD. --- src/cli/cli.ts | 31 +++++++++------ src/cli/commands/config/command.ts | 17 +++++++- src/cli/commands/fetch/command.tsx | 6 +-- src/cli/commands/pause/command.tsx | 22 +++++++---- src/cli/commands/run/command.tsx | 18 +++++++-- src/cli/commands/traces/command.tsx | 21 ++++++---- src/cli/telemetry/__tests__/client.test.ts | 38 +++++++++++++++++- src/cli/telemetry/cli-command-run.ts | 39 +++++++++++++++++-- .../schemas/__tests__/command-run.test.ts | 5 +++ src/cli/telemetry/schemas/command-run.ts | 6 +++ src/cli/telemetry/schemas/common-shapes.ts | 2 + .../fetch-access/useFetchAccessFlow.ts | 24 +++++++++--- 12 files changed, 183 insertions(+), 46 deletions(-) diff --git a/src/cli/cli.ts b/src/cli/cli.ts index 30ec862d7..194d8ca28 100644 --- a/src/cli/cli.ts +++ b/src/cli/cli.ts @@ -38,6 +38,7 @@ import { COMMAND_DESCRIPTIONS, PACKAGE_VERSION } from './constants'; import { printPostCommandNotices, printTelemetryNotice } from './notices'; import { ALL_PRIMITIVES } from './primitives'; import { TelemetryClientAccessor } from './telemetry'; +import { finalizeAndExit, registerPostCommandFinalize } from './telemetry/cli-command-run.js'; import { renderTUI, setupAltScreenCleanup } from './tui'; import { LayoutProvider } from './tui/context'; import { clearExitMessage, getExitMessage } from './tui/exit-message'; @@ -165,18 +166,24 @@ export const main = async (argv: string[]) => { } await TelemetryClientAccessor.init(args[0] ?? 'unknown'); - try { - await program.parseAsync(argv); - } finally { - await TelemetryClientAccessor.shutdown(); - } - // Telemetry notice already printed above; only run update check here. - await printPostCommandNotices(false, updateCheck); + // Post-command notices + exit message run once, whether the command returns + // normally or calls process.exit() inside its action. finalizeAndExit invokes + // this after telemetry shutdown (which prints the audit-mode line), so the + // tail output is never dropped by an early process.exit(). + registerPostCommandFinalize(async () => { + // Telemetry notice already printed above; only run update check here. + await printPostCommandNotices(false, updateCheck); + + const exitMessage = getExitMessage(); + if (exitMessage) { + console.log(`\n${exitMessage}`); + clearExitMessage(); + } + }); + + await program.parseAsync(argv); - const exitMessage = getExitMessage(); - if (exitMessage) { - console.log(`\n${exitMessage}`); - clearExitMessage(); - } + // Command returned without calling process.exit(); run the same finalize path. + await finalizeAndExit(0); }; diff --git a/src/cli/commands/config/command.ts b/src/cli/commands/config/command.ts index 8120fd806..136dd97e8 100644 --- a/src/cli/commands/config/command.ts +++ b/src/cli/commands/config/command.ts @@ -1,4 +1,6 @@ import { COMMAND_DESCRIPTIONS } from '../../constants.js'; +import { finalizeAndExit, withCommandRunTelemetry } from '../../telemetry/cli-command-run.js'; +import type { CommandAttrs } from '../../telemetry/schemas/command-run.js'; import { handleConfigGet, handleConfigList, handleConfigSet } from './actions.js'; import type { ConfigResult } from './types.js'; import type { Command } from '@commander-js/extra-typings'; @@ -9,6 +11,13 @@ function resolveAction(key?: string, value?: string): () => Promise handleConfigSet(key, value); } +// key/value are never recorded (PII/secret risk); only the derived action verb. +function deriveConfigAction(key?: string, value?: string): CommandAttrs<'config'>['config_action'] { + if (!key) return 'list'; + if (value === undefined) return 'get'; + return 'set'; +} + function printResult(result: ConfigResult): void { if (result.success) { console.log(result.message); @@ -24,8 +33,12 @@ export function registerConfig(program: Command) { .argument('[key]', 'Config key in dot notation (e.g. telemetry.enabled)') .argument('[value]', 'Value to set') .action(async (key?: string, value?: string) => { - const result = await resolveAction(key, value)(); + const result = await withCommandRunTelemetry( + 'config', + { config_action: deriveConfigAction(key, value) }, + () => resolveAction(key, value)() + ); printResult(result); - if (!result.success) process.exit(1); + await finalizeAndExit(result.success ? 0 : 1); }); } diff --git a/src/cli/commands/fetch/command.tsx b/src/cli/commands/fetch/command.tsx index dce9b47ed..812e38141 100644 --- a/src/cli/commands/fetch/command.tsx +++ b/src/cli/commands/fetch/command.tsx @@ -1,6 +1,6 @@ import { COMMAND_DESCRIPTIONS } from '../../constants'; import { getErrorMessage } from '../../errors'; -import { withCommandRunTelemetry } from '../../telemetry/cli-command-run.js'; +import { finalizeAndExit, withCommandRunTelemetry } from '../../telemetry/cli-command-run.js'; import { ResourceType, standardize } from '../../telemetry/schemas/common-shapes.js'; import { requireProject } from '../../tui/guards'; import { handleFetchAccess } from './action'; @@ -48,7 +48,7 @@ export const registerFetch = (program: Command) => { } else { render(Error: {getErrorMessage(error)}); } - process.exit(1); + await finalizeAndExit(1); return; } @@ -77,7 +77,7 @@ export const registerFetch = (program: Command) => { ); } - process.exit(1); + await finalizeAndExit(1); return; } diff --git a/src/cli/commands/pause/command.tsx b/src/cli/commands/pause/command.tsx index d790fe191..ad0a3a5b6 100644 --- a/src/cli/commands/pause/command.tsx +++ b/src/cli/commands/pause/command.tsx @@ -3,7 +3,7 @@ import { getErrorMessage } from '../../errors'; import { handlePauseResume } from '../../operations/eval'; import type { OnlineEvalActionOptions } from '../../operations/eval'; import { createJobEngine } from '../../operations/jobs'; -import { runCliCommand } from '../../telemetry/cli-command-run'; +import { finalizeAndExit, runCliCommand, withCommandRunTelemetry } from '../../telemetry/cli-command-run'; import { COMMAND_DESCRIPTIONS } from '../../tui/copy'; import { requireProject } from '../../tui/guards'; import type { Command } from '@commander-js/extra-typings'; @@ -47,7 +47,11 @@ function registerOnlineEvalSubcommand(parent: Command, action: 'pause' | 'resume }; try { - const result = await handlePauseResume(options, action); + const result = await withCommandRunTelemetry( + action === 'pause' ? 'pause.online-eval' : 'resume.online-eval', + { ref_type: cliOptions.arn ? 'arn' : 'name' }, + () => handlePauseResume(options, action) + ); if (cliOptions.json) { console.log(JSON.stringify(serializeResult(result))); @@ -58,14 +62,14 @@ function registerOnlineEvalSubcommand(parent: Command, action: 'pause' | 'resume render({result.error.message}); } - process.exit(result.success ? 0 : 1); + await finalizeAndExit(result.success ? 0 : 1); } catch (error) { if (cliOptions.json) { console.log(JSON.stringify({ success: false, error: getErrorMessage(error) })); } else { render(Error: {getErrorMessage(error)}); } - process.exit(1); + await finalizeAndExit(1); } }); } @@ -138,7 +142,11 @@ function registerOnlineInsightsSubcommand(parent: Command, action: 'pause' | 're }; try { - const result = await handlePauseResume(options, action); + const result = await withCommandRunTelemetry( + action === 'pause' ? 'pause.online-insights' : 'resume.online-insights', + { ref_type: cliOptions.arn ? 'arn' : 'name' }, + () => handlePauseResume(options, action) + ); if (cliOptions.json) { console.log(JSON.stringify(serializeResult(result))); @@ -149,14 +157,14 @@ function registerOnlineInsightsSubcommand(parent: Command, action: 'pause' | 're render({result.error.message}); } - process.exit(result.success ? 0 : 1); + await finalizeAndExit(result.success ? 0 : 1); } catch (error) { if (cliOptions.json) { console.log(JSON.stringify({ success: false, error: getErrorMessage(error) })); } else { render(Error: {getErrorMessage(error)}); } - process.exit(1); + await finalizeAndExit(1); } }); } diff --git a/src/cli/commands/run/command.tsx b/src/cli/commands/run/command.tsx index a3ce25c72..793f557f6 100644 --- a/src/cli/commands/run/command.tsx +++ b/src/cli/commands/run/command.tsx @@ -17,7 +17,7 @@ import type { StartBatchEvaluationJobOptions, } from '../../operations/jobs'; import { printABTestDetail } from '../../operations/jobs/ab-test/format'; -import { runCliCommand } from '../../telemetry/cli-command-run'; +import { finalizeAndExit, runCliCommand, withCommandRunTelemetry } from '../../telemetry/cli-command-run'; import { requireProject } from '../../tui/guards'; import type { Command } from '@commander-js/extra-typings'; import { Text, render } from 'ink'; @@ -168,7 +168,17 @@ export const registerRun = (program: Command) => { }; try { - const result = await handleRunEval(options); + const result = await withCommandRunTelemetry( + 'run.eval', + { + evaluator_count: options.evaluator.length + (options.evaluatorArn?.length ?? 0), + ref_type: options.agentArn ? 'arn' : 'name', + has_assertions: !!options.assertions?.length, + has_expected_trajectory: !!options.expectedTrajectory?.length, + has_expected_response: !!options.expectedResponse, + }, + () => handleRunEval(options) + ); if (cliOptions.json) { console.log(JSON.stringify(serializeResult(result))); @@ -179,14 +189,14 @@ export const registerRun = (program: Command) => { render({result.error.message}); } - process.exit(result.success ? 0 : 1); + await finalizeAndExit(result.success ? 0 : 1); } catch (error) { if (cliOptions.json) { console.log(JSON.stringify({ success: false, error: getErrorMessage(error) })); } else { render(Error: {getErrorMessage(error)}); } - process.exit(1); + await finalizeAndExit(1); } } ); diff --git a/src/cli/commands/traces/command.tsx b/src/cli/commands/traces/command.tsx index e15088d04..56af703bc 100644 --- a/src/cli/commands/traces/command.tsx +++ b/src/cli/commands/traces/command.tsx @@ -1,6 +1,7 @@ import { COMMAND_DESCRIPTIONS } from '../../constants'; import { getErrorMessage } from '../../errors'; import { loadDeployedProjectConfig } from '../../operations/resolve-agent'; +import { finalizeAndExit, withCommandRunTelemetry } from '../../telemetry/cli-command-run.js'; import { requireProject } from '../../tui/guards'; import { handleTracesGet, handleTracesList } from './action'; import type { TracesGetOptions, TracesListOptions } from './types'; @@ -31,8 +32,10 @@ export const registerTraces = (program: Command) => { requireProject(); try { - const context = await loadDeployedProjectConfig(); - const result = await handleTracesList(context, cliOptions); + const result = await withCommandRunTelemetry('traces.list', {}, async () => { + const context = await loadDeployedProjectConfig(); + return handleTracesList(context, cliOptions); + }); if (!result.success) { render( @@ -41,7 +44,7 @@ export const registerTraces = (program: Command) => { {result.consoleUrl && Console: {result.consoleUrl}} ); - process.exit(1); + await finalizeAndExit(1); return; } @@ -88,7 +91,7 @@ export const registerTraces = (program: Command) => { ); } catch (error) { render(Error: {getErrorMessage(error)}); - process.exit(1); + await finalizeAndExit(1); } }); @@ -103,8 +106,10 @@ export const registerTraces = (program: Command) => { requireProject(); try { - const context = await loadDeployedProjectConfig(); - const result = await handleTracesGet(context, traceId, cliOptions); + const result = await withCommandRunTelemetry('traces.get', {}, async () => { + const context = await loadDeployedProjectConfig(); + return handleTracesGet(context, traceId, cliOptions); + }); if (!result.success) { render( @@ -113,7 +118,7 @@ export const registerTraces = (program: Command) => { {result.consoleUrl && Console: {result.consoleUrl}} ); - process.exit(1); + await finalizeAndExit(1); return; } @@ -125,7 +130,7 @@ export const registerTraces = (program: Command) => { ); } catch (error) { render(Error: {getErrorMessage(error)}); - process.exit(1); + await finalizeAndExit(1); } }); }; diff --git a/src/cli/telemetry/__tests__/client.test.ts b/src/cli/telemetry/__tests__/client.test.ts index a4e280b1b..c0613ab23 100644 --- a/src/cli/telemetry/__tests__/client.test.ts +++ b/src/cli/telemetry/__tests__/client.test.ts @@ -1,9 +1,12 @@ /* eslint-disable @typescript-eslint/require-await */ import { AccessDeniedError, DependencyCheckError } from '../../../lib/errors/types'; -import { withCommandRunTelemetry } from '../cli-command-run'; +import { finalizeAndExit, registerPostCommandFinalize, withCommandRunTelemetry } from '../cli-command-run'; import { TelemetryClient } from '../client'; import { TelemetryClientAccessor } from '../client-accessor'; +import { FileSystemSink } from '../sinks/filesystem-sink'; import { InMemorySink } from '../sinks/in-memory-sink'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; let sink: InMemorySink; @@ -300,3 +303,36 @@ describe('withCommandRunTelemetry', () => { }); }); }); + +describe('finalizeAndExit', () => { + afterEach(() => { + registerPostCommandFinalize(async () => undefined); + vi.restoreAllMocks(); + }); + + it('flushes the audit sink and runs post-command notices before process.exit', async () => { + // process.exit is mocked to throw so the test runner survives the exit-path call. + const exit = vi.spyOn(process, 'exit').mockImplementation((() => { + throw new Error('exit'); + }) as never); + + const logged: string[] = []; + const auditSink = new FileSystemSink({ + filePath: join(tmpdir(), 'finalize-audit-test.jsonl'), + log: msg => logged.push(msg), + }); + auditSink.record('cli.command_run', 1, { command: 'config', exit_reason: 'success' }); + // shutdown() is the only path that emits the audit-mode line; route the accessor through it. + vi.spyOn(TelemetryClientAccessor, 'shutdown').mockImplementation(() => new TelemetryClient(auditSink).shutdown()); + + const notices = vi.fn().mockResolvedValue(undefined); + registerPostCommandFinalize(notices); + + await expect(finalizeAndExit(0)).rejects.toThrow('exit'); + + expect(logged).toHaveLength(1); + expect(logged[0]).toContain('[audit mode]'); + expect(notices).toHaveBeenCalledOnce(); + expect(exit).toHaveBeenCalledWith(0); + }); +}); diff --git a/src/cli/telemetry/cli-command-run.ts b/src/cli/telemetry/cli-command-run.ts index 6bf4818c1..c95b45916 100644 --- a/src/cli/telemetry/cli-command-run.ts +++ b/src/cli/telemetry/cli-command-run.ts @@ -11,6 +11,39 @@ import { performance } from 'perf_hooks'; export type { AttributeRecorder } from './attribute-recorder.js'; +let postCommandFinalize: (() => Promise) | undefined; + +/** + * Register work that must run after a command completes but before the process + * terminates (post-command notices, exit-message rendering). main() owns this + * because it holds the update-check promise; finalizeAndExit invokes it so that + * commands which call process.exit() inside their action still flush telemetry, + * print the audit-mode line, and show notices before exiting. + */ +export function registerPostCommandFinalize(fn: () => Promise): void { + postCommandFinalize = fn; +} + +/** + * Shut down telemetry (flushing the audit sink and printing the audit-mode + * line), run any registered post-command notices, then terminate the process. + * Centralizes the exit path so process.exit() in a command action never bypasses + * the finalize steps that main()'s finally block would otherwise run. + */ +export async function finalizeAndExit(code: number): Promise { + try { + await TelemetryClientAccessor.shutdown(); + } catch { + // Telemetry must never affect CLI behavior + } + try { + await postCommandFinalize?.(); + } catch { + // Notices are best-effort — never block exit + } + return process.exit(code); +} + async function getTelemetryClient() { try { return await TelemetryClientAccessor.get(); @@ -151,16 +184,16 @@ export async function runCliCommand( const client = await getTelemetryClient(); if (!client) { await fn(); - process.exit(0); + return finalizeAndExit(0); } await trackCommandRun(client, command, fn, knownAttrs); - process.exit(0); + return finalizeAndExit(0); } catch (error) { if (json) { console.log(JSON.stringify({ success: false, error: getErrorMessage(error) })); } else { console.error(getErrorMessage(error)); } - process.exit(1); + return finalizeAndExit(1); } } diff --git a/src/cli/telemetry/schemas/__tests__/command-run.test.ts b/src/cli/telemetry/schemas/__tests__/command-run.test.ts index d667dcdc2..940dfc860 100644 --- a/src/cli/telemetry/schemas/__tests__/command-run.test.ts +++ b/src/cli/telemetry/schemas/__tests__/command-run.test.ts @@ -127,6 +127,11 @@ describe('COMMAND_SCHEMAS', () => { expect(COMMAND_SCHEMAS['telemetry.status'].parse({})).toEqual({}); }); + it('accepts valid config attrs and rejects invalid config_action', () => { + expect(COMMAND_SCHEMAS.config.parse({ config_action: 'set' })).toEqual({ config_action: 'set' }); + expect(() => COMMAND_SCHEMAS.config.parse({ config_action: 'delete' })).toThrow(); + }); + it('import subcommand schemas accept empty object', () => { expect(COMMAND_SCHEMAS.import.parse({})).toEqual({}); expect(COMMAND_SCHEMAS['import.runtime'].parse({})).toEqual({}); diff --git a/src/cli/telemetry/schemas/command-run.ts b/src/cli/telemetry/schemas/command-run.ts index 564709f2d..e52860f6a 100644 --- a/src/cli/telemetry/schemas/command-run.ts +++ b/src/cli/telemetry/schemas/command-run.ts @@ -8,6 +8,7 @@ import { AuthType, AuthorizerType, BuildType, + ConfigAction, Count, CredentialType, DeployModeSchema, @@ -172,6 +173,10 @@ const RunIngestAttrs = safeSchema({ const FetchAccessAttrs = safeSchema({ resource_type: ResourceType }); +// config_action is derived from key/value presence. The key/value themselves are +// never recorded — they can carry secrets (e.g. telemetry endpoint) or PII. +const ConfigAttrs = safeSchema({ config_action: ConfigAction }); + /** * Async job commands (recommendation + batch-evaluation), keyed by verb with job_type to disambiguate. * safeSchema only permits required enum/boolean/number/literal fields, so per-type detail enums @@ -253,6 +258,7 @@ export const COMMAND_SCHEMAS = { 'resume.job': JobTypeOnlyAttrs, 'promote.job': JobTypeOnlyAttrs, 'fetch.access': FetchAccessAttrs, + config: ConfigAttrs, feedback: FeedbackAttrs, update: UpdateAttrs, 'pause.online-eval': PauseResumeOnlineEvalAttrs, diff --git a/src/cli/telemetry/schemas/common-shapes.ts b/src/cli/telemetry/schemas/common-shapes.ts index bf7679d80..be227c762 100644 --- a/src/cli/telemetry/schemas/common-shapes.ts +++ b/src/cli/telemetry/schemas/common-shapes.ts @@ -90,6 +90,7 @@ export const OutboundAuthType = z.enum(['oauth', 'api-key', 'none']); export const PolicyEngineMode = z.enum(['log_only', 'enforce']); export const AgentProtocol = z.enum(['http', 'mcp', 'a2a', 'agui']); export const RefType = z.enum(['arn', 'name']); +export const ConfigAction = z.enum(['get', 'list', 'set']); export const ResourceType = z.enum(['gateway', 'agent', 'harness']); export const JobType = z.enum(['recommendation', 'batch-evaluation', 'ab-test', 'insights']); export const RecommendationKind = z.enum(['system-prompt', 'tool-description']); @@ -204,6 +205,7 @@ export const ATTRIBUTES = { policy_engine_mode: PolicyEngineMode, agent_protocol: AgentProtocol, ref_type: RefType, + config_action: ConfigAction, resource_type: ResourceType, job_type: JobType, recommendation_kind: RecommendationKind, diff --git a/src/cli/tui/screens/fetch-access/useFetchAccessFlow.ts b/src/cli/tui/screens/fetch-access/useFetchAccessFlow.ts index bec80d2b8..bb8f4c6d3 100644 --- a/src/cli/tui/screens/fetch-access/useFetchAccessFlow.ts +++ b/src/cli/tui/screens/fetch-access/useFetchAccessFlow.ts @@ -9,6 +9,8 @@ import { listGateways, listHarnesses, } from '../../../operations/fetch-access'; +import { withCommandRunTelemetry } from '../../../telemetry/cli-command-run.js'; +import { ResourceType, standardize } from '../../../telemetry/schemas/common-shapes.js'; import { spawn } from 'node:child_process'; import { useCallback, useEffect, useRef, useState } from 'react'; @@ -122,12 +124,22 @@ export function useFetchAccessFlow() { const resource = state.selectedResource; - const fetchToken: Promise = - resource.resourceType === 'gateway' - ? fetchGatewayToken(resource.name) - : resource.resourceType === 'harness' - ? fetchTokenAccess(resource, fetchHarnessToken) - : fetchTokenAccess(resource, fetchRuntimeToken); + // Record cli.command_run for the TUI fetch path. The fetch result is captured via + // closure; the wrapper observes only success/throw to drive exit_reason/error_name. + let captured: TokenFetchResult; + const fetchToken: Promise = withCommandRunTelemetry( + 'fetch.access', + { resource_type: standardize(ResourceType, resource.resourceType) }, + async () => { + captured = + resource.resourceType === 'gateway' + ? await fetchGatewayToken(resource.name) + : resource.resourceType === 'harness' + ? await fetchTokenAccess(resource, fetchHarnessToken) + : await fetchTokenAccess(resource, fetchRuntimeToken); + return { success: true as const }; + } + ).then(() => captured); fetchToken .then(result => {