From f6485c65594f64c5b248dd3b81ee0fa4845ef77c Mon Sep 17 00:00:00 2001 From: Islam Shehata Date: Fri, 13 Feb 2026 10:44:39 +0200 Subject: [PATCH 1/3] feat(ai): validate token permissions before running evaluations Checks that the token has both write (ingest traces) and read (query dataset) permissions before starting evals, preventing wasted time on permission-related failures. Provides clear, actionable error messages with token management links. Co-Authored-By: Claude Sonnet 4.5 --- packages/ai/src/cli/commands/eval.command.ts | 8 + .../ai/src/config/validate-permissions.ts | 205 ++++++++++++++++++ .../test/config/validate-permissions.test.ts | 147 +++++++++++++ 3 files changed, 360 insertions(+) create mode 100644 packages/ai/src/config/validate-permissions.ts create mode 100644 packages/ai/test/config/validate-permissions.test.ts diff --git a/packages/ai/src/cli/commands/eval.command.ts b/packages/ai/src/cli/commands/eval.command.ts index 2a5c7056..c248454d 100644 --- a/packages/ai/src/cli/commands/eval.command.ts +++ b/packages/ai/src/cli/commands/eval.command.ts @@ -132,6 +132,14 @@ export const loadEvalCommand = (program: Command, flagOverrides: FlagOverrides = console.log(''); } + // Validate token permissions before running evals (skip in debug mode) + if (!options.debug) { + const { validateTokenPermissions } = await import( + '../../config/validate-permissions' + ); + await validateTokenPermissions(config); + } + const runId = createRunId(); consoleUrl = options.consoleUrl; diff --git a/packages/ai/src/config/validate-permissions.ts b/packages/ai/src/config/validate-permissions.ts new file mode 100644 index 00000000..b64843f3 --- /dev/null +++ b/packages/ai/src/config/validate-permissions.ts @@ -0,0 +1,205 @@ +import { BatchSpanProcessor, NodeTracerProvider } from '@opentelemetry/sdk-trace-node'; +import { resourceFromAttributes } from '@opentelemetry/resources'; +import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http'; +import { SpanStatusCode } from '@opentelemetry/api'; +import type { ResolvedAxiomConfig } from './index'; +import { resolveAxiomConnection } from './resolver'; +import { AxiomCLIError } from '../util/errors'; + +/** + * Result of token permission validation + */ +export interface TokenPermissionValidationResult { + valid: boolean; + canIngest: boolean; + canQuery: boolean; + errors: string[]; +} + +/** + * Validates that the configured token has the required permissions to run evaluations. + * + * This function: + * 1. Attempts to send a test trace to verify ingestion permissions + * 2. Attempts to query the dataset to verify read permissions + * + * @param config - Resolved Axiom configuration + * @returns Validation result with permission details + * @throws {AxiomCLIError} If validation fails with detailed error messages + */ +export async function validateTokenPermissions( + config: ResolvedAxiomConfig, +): Promise { + const result: TokenPermissionValidationResult = { + valid: true, + canIngest: false, + canQuery: false, + errors: [], + }; + + const connection = resolveAxiomConnection(config); + + // Test 1: Verify token can ingest traces + try { + await testTraceIngestion(connection); + result.canIngest = true; + } catch (error) { + result.valid = false; + result.canIngest = false; + const errorMsg = error instanceof Error ? error.message : 'Unknown error'; + result.errors.push(`❌ Cannot ingest traces: ${errorMsg}`); + } + + // Test 2: Verify token can query dataset + try { + await testDatasetQuery(connection); + result.canQuery = true; + } catch (error) { + result.valid = false; + result.canQuery = false; + const errorMsg = error instanceof Error ? error.message : 'Unknown error'; + result.errors.push(`❌ Cannot query dataset: ${errorMsg}`); + } + + if (!result.valid) { + const errorMessage = [ + `Token validation failed for dataset '${connection.dataset}'.`, + 'Your token lacks required permissions:', + '', + ...result.errors, + '', + 'To run evaluations, your token needs:', + ' • Write permission to ingest traces', + ' • Read permission to query results', + '', + 'Manage tokens at: ' + + connection.consoleEndpointUrl.replace('/app/', '/app/settings/tokens'), + ].join('\n'); + + throw new AxiomCLIError(errorMessage); + } + + return result; +} + +/** + * Test trace ingestion by sending a validation trace + */ +async function testTraceIngestion( + connection: ReturnType, +): Promise { + const headers: Record = { + 'X-Axiom-Dataset': connection.dataset, + ...(connection.orgId ? { 'X-AXIOM-ORG-ID': connection.orgId } : {}), + }; + + if (connection.token) { + headers.Authorization = `Bearer ${connection.token}`; + } + + // Use edgeUrl for trace ingestion + const collectorOptions = { + url: `${connection.edgeUrl}/v1/traces`, + headers, + concurrencyLimit: 1, + }; + + const exporter = new OTLPTraceExporter(collectorOptions); + + const processor = new BatchSpanProcessor(exporter, { + maxQueueSize: 10, + maxExportBatchSize: 10, + scheduledDelayMillis: 100, + exportTimeoutMillis: 5000, + }); + + const provider = new NodeTracerProvider({ + resource: resourceFromAttributes({ + ['service.name']: 'axiom-token-validation', + ['service.version']: __SDK_VERSION__, + }), + spanProcessors: [processor], + }); + + const tracer = provider.getTracer('axiom-token-validation', __SDK_VERSION__); + + // Create a validation span + const span = tracer.startSpan('axiom.token.validation', { + attributes: { + 'axiom.validation.type': 'token_permission_check', + 'axiom.validation.timestamp': Date.now(), + }, + }); + + span.setStatus({ code: SpanStatusCode.OK }); + span.end(); + + // Force flush to send the span immediately + try { + await provider.forceFlush(); + } catch (error) { + // Extract meaningful error message + let message = 'Write failed'; + if (error instanceof Error) { + // Check for common HTTP error patterns + if (error.message.includes('401') || error.message.includes('Unauthorized')) { + message = 'Invalid or expired token'; + } else if (error.message.includes('403') || error.message.includes('Forbidden')) { + message = 'Write permission denied'; + } else if (error.message.includes('404')) { + message = 'Dataset not found'; + } else { + message = error.message; + } + } + throw new Error(message); + } finally { + await provider.shutdown(); + } +} + +/** + * Test dataset query by running a simple APL query + */ +async function testDatasetQuery( + connection: ReturnType, +): Promise { + // Simple APL query that returns minimal data + const apl = `['${connection.dataset}'] | limit 1`; + + const headers = new Headers({ + Authorization: `Bearer ${connection.token}`, + 'Content-Type': 'application/json', + ...(connection.orgId ? { 'X-AXIOM-ORG-ID': connection.orgId } : {}), + }); + + try { + // Use edgeUrl for query operations + const response = await fetch(`${connection.edgeUrl}/v1/datasets/_apl?format=legacy`, { + headers, + method: 'POST', + body: JSON.stringify({ apl }), + }); + + if (!response.ok) { + const payload = await response.json().catch(() => ({})); + const errorMessage = payload.message || response.statusText; + + // Provide specific error messages based on status code + if (response.status === 401) { + throw new Error('Invalid or expired token'); + } else if (response.status === 403) { + throw new Error('Read permission denied'); + } else if (response.status === 404) { + throw new Error('Dataset not found'); + } else { + throw new Error(`Query failed: ${errorMessage}`); + } + } + } catch (error) { + if (error instanceof Error) { + throw error; + } + throw new Error('Failed to query dataset: Network error'); + } +} diff --git a/packages/ai/test/config/validate-permissions.test.ts b/packages/ai/test/config/validate-permissions.test.ts new file mode 100644 index 00000000..e82b84c1 --- /dev/null +++ b/packages/ai/test/config/validate-permissions.test.ts @@ -0,0 +1,147 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { validateTokenPermissions } from '../../src/config/validate-permissions'; +import type { ResolvedAxiomConfig } from '../../src/config'; +import { AxiomCLIError } from '../../src/util/errors'; + +// Mock the OTEL modules +vi.mock('@opentelemetry/sdk-trace-node', () => ({ + BatchSpanProcessor: vi.fn(function (this: any) { + this.forceFlush = vi.fn().mockResolvedValue(undefined); + this.shutdown = vi.fn().mockResolvedValue(undefined); + }), + NodeTracerProvider: vi.fn(function (this: any) { + this.getTracer = vi.fn().mockReturnValue({ + startSpan: vi.fn().mockReturnValue({ + setStatus: vi.fn(), + end: vi.fn(), + }), + }); + this.forceFlush = vi.fn().mockResolvedValue(undefined); + this.shutdown = vi.fn().mockResolvedValue(undefined); + }), +})); + +vi.mock('@opentelemetry/resources', () => ({ + resourceFromAttributes: vi.fn().mockReturnValue({}), +})); + +vi.mock('@opentelemetry/exporter-trace-otlp-http', () => ({ + OTLPTraceExporter: vi.fn(function (this: any) { + // Mock OTLP exporter constructor + }), +})); + +describe('validateTokenPermissions', () => { + const mockConfig: ResolvedAxiomConfig = { + eval: { + url: 'https://api.axiom.co', + edgeUrl: 'https://api.axiom.co', + token: 'test-token', + dataset: 'test-dataset', + orgId: 'test-org', + flagSchema: null, + instrumentation: null, + timeoutMs: 60000, + include: ['**/*.eval.ts'], + exclude: [], + }, + }; + + let originalFetch: typeof global.fetch; + + beforeEach(() => { + originalFetch = global.fetch; + (globalThis as any).__SDK_VERSION__ = 'test-version'; + }); + + afterEach(() => { + global.fetch = originalFetch; + vi.clearAllMocks(); + }); + + it('should pass validation when token has all permissions', async () => { + // Mock successful query response + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ matches: [] }), + }); + + const result = await validateTokenPermissions(mockConfig); + + expect(result.valid).toBe(true); + expect(result.canIngest).toBe(true); + expect(result.canQuery).toBe(true); + expect(result.errors).toHaveLength(0); + }); + + it('should fail validation when query fails with 401', async () => { + // Mock 401 unauthorized query response + global.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 401, + statusText: 'Unauthorized', + json: async () => ({ message: 'Unauthorized' }), + }); + + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(AxiomCLIError); + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Invalid or expired token/); + }); + + it('should fail validation when query fails with 403', async () => { + // Mock 403 forbidden query response + global.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 403, + statusText: 'Forbidden', + json: async () => ({ message: 'Forbidden' }), + }); + + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(AxiomCLIError); + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Read permission denied/); + }); + + it('should fail validation when query fails with 404', async () => { + // Mock 404 not found query response + global.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 404, + statusText: 'Not Found', + json: async () => ({ message: 'Not Found' }), + }); + + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(AxiomCLIError); + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Dataset not found/); + }); + + it('should provide helpful error message with token management URL', async () => { + // Mock failed query + global.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 403, + statusText: 'Forbidden', + json: async () => ({ message: 'Forbidden' }), + }); + + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Manage tokens at:/); + }); + + it('should list required permissions in error message', async () => { + // Mock failed query + global.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 403, + statusText: 'Forbidden', + json: async () => ({ message: 'Forbidden' }), + }); + + try { + await validateTokenPermissions(mockConfig); + } catch (error) { + if (error instanceof AxiomCLIError) { + expect(error.message).toContain('To run evaluations, your token needs:'); + expect(error.message).toContain('Write permission to ingest traces'); + expect(error.message).toContain('Read permission to query results'); + } + } + }); +}); From 67e5d0181fc55a858786ad24ddc5361e52fbb356 Mon Sep 17 00:00:00 2001 From: Islam Shehata Date: Wed, 18 Feb 2026 17:20:21 +0200 Subject: [PATCH 2/3] fix(evals): improve token validation errors and reporting --- examples/kitchen-sink/axiom.config.ts | 1 + packages/ai/src/cli/commands/eval.command.ts | 13 +- packages/ai/src/config/resolver.ts | 4 + .../config/validate-eval-token-permissions.ts | 143 ++++++++++++ .../ai/src/config/validate-permissions.ts | 205 ------------------ packages/ai/src/evals/eval.service.ts | 4 +- packages/ai/src/evals/eval.ts | 66 +++--- packages/ai/src/evals/eval.types.ts | 2 + .../ai/src/evals/reporter.console-utils.ts | 16 ++ packages/ai/src/evals/reporter.ts | 1 + packages/ai/test/cli/eval.command.test.ts | 95 ++++++++ ...> validate-eval-token-permissions.test.ts} | 44 ++-- 12 files changed, 336 insertions(+), 258 deletions(-) create mode 100644 packages/ai/src/config/validate-eval-token-permissions.ts delete mode 100644 packages/ai/src/config/validate-permissions.ts create mode 100644 packages/ai/test/cli/eval.command.test.ts rename packages/ai/test/config/{validate-permissions.test.ts => validate-eval-token-permissions.test.ts} (83%) diff --git a/examples/kitchen-sink/axiom.config.ts b/examples/kitchen-sink/axiom.config.ts index 4dd64489..a85a17b7 100644 --- a/examples/kitchen-sink/axiom.config.ts +++ b/examples/kitchen-sink/axiom.config.ts @@ -9,6 +9,7 @@ export default defineConfig({ dataset: process.env.NEXT_PUBLIC_AXIOM_DATASET, url: process.env.NEXT_PUBLIC_AXIOM_URL, + edgeUrl: process.env.NEXT_PUBLIC_AXIOM_EDGE_URL, token: process.env.AXIOM_TOKEN, flagSchema, diff --git a/packages/ai/src/cli/commands/eval.command.ts b/packages/ai/src/cli/commands/eval.command.ts index 80a3505b..14093059 100644 --- a/packages/ai/src/cli/commands/eval.command.ts +++ b/packages/ai/src/cli/commands/eval.command.ts @@ -1,13 +1,15 @@ import { Command, Argument, Option } from 'commander'; import { customAlphabet } from 'nanoid'; import { lstatSync } from 'node:fs'; +import c from 'tinyrainbow'; + import { runEvalWithContext } from '../utils/eval-context-runner'; import { validateFlagOverrides, type FlagOverrides } from '../utils/parse-flag-overrides'; import { isGlob } from '../utils/glob-utils'; import { loadConfig } from '../../config/loader'; import { AxiomCLIError } from '../../util/errors'; import { getAuthContext } from '../auth/global-auth'; -import c from 'tinyrainbow'; +import { validateTokenPermissions } from '../../config/validate-eval-token-permissions'; const createRunId = customAlphabet('0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ', 10); @@ -134,10 +136,11 @@ export const loadEvalCommand = (program: Command, flagOverrides: FlagOverrides = // Validate token permissions before running evals (skip in debug mode) if (!options.debug) { - const { validateTokenPermissions } = await import( - '../../config/validate-permissions' - ); - await validateTokenPermissions(config); + const result = await validateTokenPermissions(config); + if (!result.valid) { + console.error(`\n❌ ${result.errors.join('\n')} \n`); + process.exit(1); + } } const runId = createRunId(); diff --git a/packages/ai/src/config/resolver.ts b/packages/ai/src/config/resolver.ts index d523af2c..9c8ffd65 100644 --- a/packages/ai/src/config/resolver.ts +++ b/packages/ai/src/config/resolver.ts @@ -8,6 +8,10 @@ import type { AxiomEvalInstrumentationOptions, ResolvedAxiomConfig } from './ind const buildConsoleUrl = (urlString: string) => { const url = new URL(urlString); + if (url.host.startsWith('localhost:')) { + return urlString; + } + return `${url.protocol}//app.${url.host.split('api.').at(-1)}`; }; diff --git a/packages/ai/src/config/validate-eval-token-permissions.ts b/packages/ai/src/config/validate-eval-token-permissions.ts new file mode 100644 index 00000000..42c98625 --- /dev/null +++ b/packages/ai/src/config/validate-eval-token-permissions.ts @@ -0,0 +1,143 @@ +import type { ResolvedAxiomConfig } from './index'; +import { resolveAxiomConnection } from './resolver'; +import { AxiomCLIError } from '../util/errors'; + +/** + * Result of token permission validation + */ +export interface TokenPermissionValidationResult { + valid: boolean; + dataset: string; + permissions: { + canRead: boolean; + canWrite: boolean; + }; + errors: string[]; +} + +/** + * Validates that the configured token has the required permissions to run evaluations. + * + * This function: + * 1. Attempts to send a test trace to verify ingestion permissions + * 2. Attempts to query the dataset to verify read permissions + * + * @param config - Resolved Axiom configuration + * @returns Validation result with permission details + * @throws {AxiomCLIError} If validation fails with detailed error messages + */ +const buildPermissionHelp = (consoleEndpointUrl: string) => [ + 'To run evaluations, your token needs:', + ' - Write permission to ingest traces', + ' - Read permission to query results', + `Manage tokens at: ${consoleEndpointUrl}/settings/api-tokens`, +]; + +const indentErrorDetails = (lines: string[]) => + lines.map((line, index) => (index === 0 ? line : ` ${line}`)); + +const formatPermissionErrors = ( + errors: string[] | undefined, + dataset: string, + consoleEndpointUrl: string, +) => { + const details: string[] = []; + const normalized = (errors || []).map((e) => e.toLowerCase()); + + if (normalized.some((e) => e.includes('ingest') || e.includes('write'))) { + details.push('Missing write permission to ingest traces.'); + } + if (normalized.some((e) => e.includes('read') || e.includes('query'))) { + details.push('Missing read permission to query results.'); + } + + if (details.length === 0 && errors && errors.length > 0) { + details.push(...errors); + } + + return indentErrorDetails([ + `Token does not have required permissions for dataset "${dataset}".`, + ...details, + ...buildPermissionHelp(consoleEndpointUrl), + ]); +}; + +export async function validateTokenPermissions(config: ResolvedAxiomConfig) { + const connection = resolveAxiomConnection(config); + + try { + const headers: Record = { + 'X-Axiom-Dataset': connection.dataset, + }; + + if (connection.token) { + headers.Authorization = `Bearer ${connection.token}`; + } + + const edge = new URL(connection.edgeUrl); + const region = edge.hostname || 'us-east-1'; + console.log({ region }); + const response = await fetch( + `${connection.url}/api/v3/evaluations/validate?dataset=${connection.dataset}®ion=${region}`, + { + headers, + }, + ); + + if (!response.ok) { + let serverMessage: string | undefined; + try { + const data = await response.json(); + serverMessage = data?.error || data?.message; + } catch { + serverMessage = undefined; + } + + if (response.status === 404) { + throw new AxiomCLIError( + indentErrorDetails([ + `Dataset not found: "${connection.dataset}".`, + 'Check eval.dataset in axiom.config.ts or AXIOM_DATASET in your environment.', + `Manage datasets at: ${connection.consoleEndpointUrl}/datasets`, + ]).join('\n'), + ); + } + + const statusLabel = serverMessage || response.statusText || 'Unknown error'; + const baseMessage = `Failed to validate token: ${statusLabel}`; + + if (response.status === 401 || response.status === 403) { + throw new AxiomCLIError( + indentErrorDetails([ + baseMessage, + response.status === 401 + ? 'The token is missing or invalid.' + : `The token does not have access to dataset "${connection.dataset}".`, + `Check AXIOM_TOKEN or eval.token in axiom.config.ts.`, + ...buildPermissionHelp(connection.consoleEndpointUrl), + ]).join('\n'), + ); + } + + throw new AxiomCLIError( + indentErrorDetails([ + baseMessage, + ...buildPermissionHelp(connection.consoleEndpointUrl), + ]).join('\n'), + ); + } + + const result = (await response.json()) as TokenPermissionValidationResult; + if (!result.valid) { + result.errors = formatPermissionErrors( + result.errors, + connection.dataset, + connection.consoleEndpointUrl, + ); + } + return result; + } catch (error) { + const errorMessage = error instanceof Error ? error.message : 'Unknown error'; + throw new AxiomCLIError(errorMessage); + } +} diff --git a/packages/ai/src/config/validate-permissions.ts b/packages/ai/src/config/validate-permissions.ts deleted file mode 100644 index b64843f3..00000000 --- a/packages/ai/src/config/validate-permissions.ts +++ /dev/null @@ -1,205 +0,0 @@ -import { BatchSpanProcessor, NodeTracerProvider } from '@opentelemetry/sdk-trace-node'; -import { resourceFromAttributes } from '@opentelemetry/resources'; -import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http'; -import { SpanStatusCode } from '@opentelemetry/api'; -import type { ResolvedAxiomConfig } from './index'; -import { resolveAxiomConnection } from './resolver'; -import { AxiomCLIError } from '../util/errors'; - -/** - * Result of token permission validation - */ -export interface TokenPermissionValidationResult { - valid: boolean; - canIngest: boolean; - canQuery: boolean; - errors: string[]; -} - -/** - * Validates that the configured token has the required permissions to run evaluations. - * - * This function: - * 1. Attempts to send a test trace to verify ingestion permissions - * 2. Attempts to query the dataset to verify read permissions - * - * @param config - Resolved Axiom configuration - * @returns Validation result with permission details - * @throws {AxiomCLIError} If validation fails with detailed error messages - */ -export async function validateTokenPermissions( - config: ResolvedAxiomConfig, -): Promise { - const result: TokenPermissionValidationResult = { - valid: true, - canIngest: false, - canQuery: false, - errors: [], - }; - - const connection = resolveAxiomConnection(config); - - // Test 1: Verify token can ingest traces - try { - await testTraceIngestion(connection); - result.canIngest = true; - } catch (error) { - result.valid = false; - result.canIngest = false; - const errorMsg = error instanceof Error ? error.message : 'Unknown error'; - result.errors.push(`❌ Cannot ingest traces: ${errorMsg}`); - } - - // Test 2: Verify token can query dataset - try { - await testDatasetQuery(connection); - result.canQuery = true; - } catch (error) { - result.valid = false; - result.canQuery = false; - const errorMsg = error instanceof Error ? error.message : 'Unknown error'; - result.errors.push(`❌ Cannot query dataset: ${errorMsg}`); - } - - if (!result.valid) { - const errorMessage = [ - `Token validation failed for dataset '${connection.dataset}'.`, - 'Your token lacks required permissions:', - '', - ...result.errors, - '', - 'To run evaluations, your token needs:', - ' • Write permission to ingest traces', - ' • Read permission to query results', - '', - 'Manage tokens at: ' + - connection.consoleEndpointUrl.replace('/app/', '/app/settings/tokens'), - ].join('\n'); - - throw new AxiomCLIError(errorMessage); - } - - return result; -} - -/** - * Test trace ingestion by sending a validation trace - */ -async function testTraceIngestion( - connection: ReturnType, -): Promise { - const headers: Record = { - 'X-Axiom-Dataset': connection.dataset, - ...(connection.orgId ? { 'X-AXIOM-ORG-ID': connection.orgId } : {}), - }; - - if (connection.token) { - headers.Authorization = `Bearer ${connection.token}`; - } - - // Use edgeUrl for trace ingestion - const collectorOptions = { - url: `${connection.edgeUrl}/v1/traces`, - headers, - concurrencyLimit: 1, - }; - - const exporter = new OTLPTraceExporter(collectorOptions); - - const processor = new BatchSpanProcessor(exporter, { - maxQueueSize: 10, - maxExportBatchSize: 10, - scheduledDelayMillis: 100, - exportTimeoutMillis: 5000, - }); - - const provider = new NodeTracerProvider({ - resource: resourceFromAttributes({ - ['service.name']: 'axiom-token-validation', - ['service.version']: __SDK_VERSION__, - }), - spanProcessors: [processor], - }); - - const tracer = provider.getTracer('axiom-token-validation', __SDK_VERSION__); - - // Create a validation span - const span = tracer.startSpan('axiom.token.validation', { - attributes: { - 'axiom.validation.type': 'token_permission_check', - 'axiom.validation.timestamp': Date.now(), - }, - }); - - span.setStatus({ code: SpanStatusCode.OK }); - span.end(); - - // Force flush to send the span immediately - try { - await provider.forceFlush(); - } catch (error) { - // Extract meaningful error message - let message = 'Write failed'; - if (error instanceof Error) { - // Check for common HTTP error patterns - if (error.message.includes('401') || error.message.includes('Unauthorized')) { - message = 'Invalid or expired token'; - } else if (error.message.includes('403') || error.message.includes('Forbidden')) { - message = 'Write permission denied'; - } else if (error.message.includes('404')) { - message = 'Dataset not found'; - } else { - message = error.message; - } - } - throw new Error(message); - } finally { - await provider.shutdown(); - } -} - -/** - * Test dataset query by running a simple APL query - */ -async function testDatasetQuery( - connection: ReturnType, -): Promise { - // Simple APL query that returns minimal data - const apl = `['${connection.dataset}'] | limit 1`; - - const headers = new Headers({ - Authorization: `Bearer ${connection.token}`, - 'Content-Type': 'application/json', - ...(connection.orgId ? { 'X-AXIOM-ORG-ID': connection.orgId } : {}), - }); - - try { - // Use edgeUrl for query operations - const response = await fetch(`${connection.edgeUrl}/v1/datasets/_apl?format=legacy`, { - headers, - method: 'POST', - body: JSON.stringify({ apl }), - }); - - if (!response.ok) { - const payload = await response.json().catch(() => ({})); - const errorMessage = payload.message || response.statusText; - - // Provide specific error messages based on status code - if (response.status === 401) { - throw new Error('Invalid or expired token'); - } else if (response.status === 403) { - throw new Error('Read permission denied'); - } else if (response.status === 404) { - throw new Error('Dataset not found'); - } else { - throw new Error(`Query failed: ${errorMessage}`); - } - } - } catch (error) { - if (error instanceof Error) { - throw error; - } - throw new Error('Failed to query dataset: Network error'); - } -} diff --git a/packages/ai/src/evals/eval.service.ts b/packages/ai/src/evals/eval.service.ts index 862e680d..b163acec 100644 --- a/packages/ai/src/evals/eval.service.ts +++ b/packages/ai/src/evals/eval.service.ts @@ -3,7 +3,7 @@ import { createFetcher, type Fetcher } from '../utils/fetcher'; import type { ResolvedAxiomConfig } from '../config/index'; import { resolveAxiomConnection } from '../config/resolver'; import { Attr } from '../otel'; -import { AxiomCLIError } from '../util/errors'; +import { AxiomCLIError, errorToString } from '../util/errors'; import { getCustomOrRegularAttribute, getCustomOrRegularNumber, @@ -52,6 +52,8 @@ export class EvaluationApiClient { const resp = await this.fetcher(`/api/v3/evaluations`, { method: 'POST', body: JSON.stringify(evaluation), + }).catch((error) => { + throw new AxiomCLIError(`Failed to create evaluation: ${errorToString(error)}`); }); if (!resp.ok) { diff --git a/packages/ai/src/evals/eval.ts b/packages/ai/src/evals/eval.ts index f53cbe18..817cb361 100644 --- a/packages/ai/src/evals/eval.ts +++ b/packages/ai/src/evals/eval.ts @@ -258,6 +258,7 @@ async function registerEval< const [, instrumentationInitError] = await tryCatchAsync(instrumentationReady); if (instrumentationInitError) { instrumentationError = instrumentationInitError; + suite.meta.evaluation.instrumentationError = errorToString(instrumentationInitError); } suiteSpan = startSpan(`eval ${evalName}-${evalVersion}`, { @@ -293,22 +294,28 @@ async function registerEval< suiteSpan.setAttribute(Attr.Eval.Config.Flags, flagConfigJson); let createEvalResponse; + let registrationError: Error | undefined = undefined; if (!isDebug && !isList) { - createEvalResponse = await evaluationApiClient.createEvaluation({ - id: evalId, - name: evalName, - capability: opts.capability, - step: opts.step, - dataset: axiomConfig.eval.dataset, - version: evalVersion, - baselineId: baselineId ?? undefined, - runId: runId, - totalCases: collection.length, - config: { overrides: injectedOverrides }, - configTimeoutMs: timeoutMs, - metadata: opts.metadata, - status: 'running', - }); + try { + createEvalResponse = await evaluationApiClient.createEvaluation({ + id: evalId, + name: evalName, + capability: opts.capability, + step: opts.step, + dataset: axiomConfig.eval.dataset, + version: evalVersion, + baselineId: baselineId ?? undefined, + runId: runId, + totalCases: collection.length, + config: { overrides: injectedOverrides }, + configTimeoutMs: timeoutMs, + metadata: opts.metadata, + status: 'running', + }); + registrationError = createEvalResponse?.error; + } catch (error) { + registrationError = error as Error; + } } const orgId = createEvalResponse?.data?.orgId; @@ -342,10 +349,10 @@ async function registerEval< orgId: orgId ?? undefined, baseline: baseline ?? undefined, configFlags: opts.configFlags, - registrationStatus: instrumentationError + registrationStatus: registrationError ? { status: 'failed', - error: errorToString(instrumentationError), + error: errorToString(registrationError), } : { status: 'success' }, trials: opts.trials, @@ -424,14 +431,23 @@ async function registerEval< // signal Axiom that evaluation finished to kick of summary calculations if (!isDebug && !isList) { - await evaluationApiClient.updateEvaluation({ - id: evalId, - status: 'completed', - totalCases: collection.length, - successCases, - erroredCases, - durationMs, - }); + try { + await evaluationApiClient.updateEvaluation({ + id: evalId, + status: 'completed', + totalCases: collection.length, + successCases, + erroredCases, + durationMs, + }); + } catch (error) { + if (suite.meta.evaluation?.registrationStatus?.status !== 'failed') { + suite.meta.evaluation.registrationStatus = { + status: 'failed', + error: errorToString(error), + }; + } + } } }); diff --git a/packages/ai/src/evals/eval.types.ts b/packages/ai/src/evals/eval.types.ts index a23ac567..729e9f31 100644 --- a/packages/ai/src/evals/eval.types.ts +++ b/packages/ai/src/evals/eval.types.ts @@ -265,6 +265,8 @@ export type EvaluationReport = { overrides?: Record; }; registrationStatus?: RegistrationStatus; + /** Captures instrumentation/baseline loading failures for reporting */ + instrumentationError?: string; /** Number of trials per case (only shown if > 1) */ trials?: number; }; diff --git a/packages/ai/src/evals/reporter.console-utils.ts b/packages/ai/src/evals/reporter.console-utils.ts index 40bf3013..0e5de723 100644 --- a/packages/ai/src/evals/reporter.console-utils.ts +++ b/packages/ai/src/evals/reporter.console-utils.ts @@ -41,6 +41,7 @@ export type SuiteData = { }>; outOfScopeFlags?: OutOfScopeFlag[]; registrationStatus?: RegistrationStatus; + instrumentationError?: string; }; export type Logger = (message?: string, ...optionalParams: any[]) => void; @@ -601,6 +602,7 @@ export function printFinalReport({ logger(''); } } + const anyInstrumentationErrors = suiteData.some((s) => !!s.instrumentationError); if (anyRegistered && orgId && config?.consoleEndpointUrl) { if (suiteData.length === 1) { @@ -632,4 +634,18 @@ export function printFinalReport({ } } } + + if (anyInstrumentationErrors) { + logger(''); + for (const suite of suiteData) { + if (!suite.instrumentationError) continue; + logger(c.yellow(`⚠️ Warning: Instrumentation error in "${suite.name}"`)); + logger(c.dim(` Error: ${suite.instrumentationError}`)); + logger( + c.dim( + ' Some instrumentation or baseline features may be unavailable for this evaluation.', + ), + ); + } + } } diff --git a/packages/ai/src/evals/reporter.ts b/packages/ai/src/evals/reporter.ts index b760b39e..7e8b4c01 100644 --- a/packages/ai/src/evals/reporter.ts +++ b/packages/ai/src/evals/reporter.ts @@ -137,6 +137,7 @@ export class AxiomReporter implements Reporter { cases, outOfScopeFlags: meta.evaluation.outOfScopeFlags, registrationStatus: meta.evaluation.registrationStatus, + instrumentationError: meta.evaluation.instrumentationError, }); printEvalNameAndFileName(testSuite, meta); diff --git a/packages/ai/test/cli/eval.command.test.ts b/packages/ai/test/cli/eval.command.test.ts new file mode 100644 index 00000000..b4269ae5 --- /dev/null +++ b/packages/ai/test/cli/eval.command.test.ts @@ -0,0 +1,95 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { Command } from 'commander'; + +import { loadEvalCommand } from '../../src/cli/commands/eval.command'; + +vi.mock('../../src/config/loader', () => ({ + loadConfig: vi.fn().mockResolvedValue({ + config: { + eval: { + url: 'https://api.axiom.co', + edgeUrl: 'https://api.axiom.co', + token: 'test-token', + dataset: 'test-dataset', + orgId: 'test-org', + flagSchema: null, + instrumentation: {}, + timeoutMs: 60000, + include: ['**/*.eval.ts'], + exclude: [], + }, + }, + }), +})); + +vi.mock('../../src/cli/utils/eval-context-runner', () => ({ + runEvalWithContext: vi.fn(async (_overrides: unknown, fn: () => Promise) => fn()), +})); + +vi.mock('../../src/evals/run-vitest', () => ({ + runVitest: vi.fn().mockResolvedValue(undefined), +})); + +describe('eval command token validation', () => { + let originalFetch: typeof global.fetch; + let exitSpy: ReturnType; + let consoleErrorSpy: ReturnType; + + beforeEach(() => { + originalFetch = global.fetch; + (globalThis as any).__SDK_VERSION__ = 'test-version'; + exitSpy = vi.spyOn(process, 'exit').mockImplementation(((code?: number) => { + throw new Error(`process.exit: ${code}`); + }) as any); + consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}); + }); + + afterEach(() => { + global.fetch = originalFetch; + exitSpy.mockRestore(); + consoleErrorSpy.mockRestore(); + vi.clearAllMocks(); + }); + + it('prints validation errors returned from the validate endpoint', async () => { + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + valid: false, + permissions: { + canWrite: false, + canRead: false, + }, + errors: [ + 'the token does not have ingest capability for the dataset', + 'the token does not have read capability for the dataset', + ], + }), + }); + + const program = new Command(); + loadEvalCommand(program); + + await expect( + program.parseAsync([ + 'node', + 'axiom', + 'eval', + '--token', + 'test-token', + '--dataset', + 'test-dataset', + '--url', + 'https://api.axiom.co', + '--org-id', + 'test-org', + ]), + ).rejects.toThrow(/process\.exit: 1/); + + expect(consoleErrorSpy).toHaveBeenCalled(); + const errorOutput = consoleErrorSpy.mock.calls.map((call) => call.join(' ')).join('\n'); + expect(errorOutput).toContain('Token does not have required permissions for dataset "test-dataset".'); + expect(errorOutput).toContain('Missing write permission to ingest traces.'); + expect(errorOutput).toContain('Missing read permission to query results.'); + }); +}); diff --git a/packages/ai/test/config/validate-permissions.test.ts b/packages/ai/test/config/validate-eval-token-permissions.test.ts similarity index 83% rename from packages/ai/test/config/validate-permissions.test.ts rename to packages/ai/test/config/validate-eval-token-permissions.test.ts index e82b84c1..a85275b4 100644 --- a/packages/ai/test/config/validate-permissions.test.ts +++ b/packages/ai/test/config/validate-eval-token-permissions.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; -import { validateTokenPermissions } from '../../src/config/validate-permissions'; -import type { ResolvedAxiomConfig } from '../../src/config'; +import { validateTokenPermissions } from '../../src/config/validate-eval-token-permissions'; +import type { ResolvedAxiomConfig } from '../../src/config/index'; import { AxiomCLIError } from '../../src/util/errors'; // Mock the OTEL modules @@ -63,14 +63,18 @@ describe('validateTokenPermissions', () => { // Mock successful query response global.fetch = vi.fn().mockResolvedValue({ ok: true, - json: async () => ({ matches: [] }), + json: async () => ({ + valid: true, + permissions: { canWrite: true, canRead: true }, + errors: [], + }), }); const result = await validateTokenPermissions(mockConfig); expect(result.valid).toBe(true); - expect(result.canIngest).toBe(true); - expect(result.canQuery).toBe(true); + expect(result.permissions.canWrite).toBe(true); + expect(result.permissions.canRead).toBe(true); expect(result.errors).toHaveLength(0); }); @@ -84,7 +88,9 @@ describe('validateTokenPermissions', () => { }); await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(AxiomCLIError); - await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Invalid or expired token/); + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow( + /Failed to validate token: Unauthorized/, + ); }); it('should fail validation when query fails with 403', async () => { @@ -97,7 +103,9 @@ describe('validateTokenPermissions', () => { }); await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(AxiomCLIError); - await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Read permission denied/); + await expect(validateTokenPermissions(mockConfig)).rejects.toThrow( + /Failed to validate token: Forbidden/, + ); }); it('should fail validation when query fails with 404', async () => { @@ -113,25 +121,17 @@ describe('validateTokenPermissions', () => { await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Dataset not found/); }); - it('should provide helpful error message with token management URL', async () => { - // Mock failed query - global.fetch = vi.fn().mockResolvedValue({ - ok: false, - status: 403, - statusText: 'Forbidden', - json: async () => ({ message: 'Forbidden' }), - }); - - await expect(validateTokenPermissions(mockConfig)).rejects.toThrow(/Manage tokens at:/); - }); - it('should list required permissions in error message', async () => { // Mock failed query global.fetch = vi.fn().mockResolvedValue({ ok: false, - status: 403, - statusText: 'Forbidden', - json: async () => ({ message: 'Forbidden' }), + status: 200, + statusText: '', + json: async () => ({ + valid: false, + permissions: { canWrite: false, canRead: false }, + errors: [], + }), }); try { From e126238a459cf6a37e6c2e2d384a08eac86d1b67 Mon Sep 17 00:00:00 2001 From: Islam Shehata Date: Thu, 12 Mar 2026 01:25:06 +0200 Subject: [PATCH 3/3] save progress --- packages/ai/src/config/resolver.ts | 25 +++++++++++++++- .../config/validate-eval-token-permissions.ts | 9 +++--- packages/ai/test/config/load-config.test.ts | 30 +++++++++++++++++++ 3 files changed, 59 insertions(+), 5 deletions(-) diff --git a/packages/ai/src/config/resolver.ts b/packages/ai/src/config/resolver.ts index 9c8ffd65..795bb9c4 100644 --- a/packages/ai/src/config/resolver.ts +++ b/packages/ai/src/config/resolver.ts @@ -1,5 +1,26 @@ import type { AxiomEvalInstrumentationOptions, ResolvedAxiomConfig } from './index'; +const DEFAULT_EDGE_REGION = 'us-east-1'; +const NON_EDGE_HOSTS = new Set(['api.axiom.co', 'api.dev.axiomtestlabs.co']); +const LOCALHOST_HOSTS = new Set(['localhost', '127.0.0.1', '::1', '0.0.0.0']); + +export function resolveEdgeRegion(edgeUrl: string): string { + let hostname = ''; + + try { + hostname = new URL(edgeUrl).hostname.toLowerCase(); + } catch { + return DEFAULT_EDGE_REGION; + } + + if (NON_EDGE_HOSTS.has(hostname) || LOCALHOST_HOSTS.has(hostname)) { + return DEFAULT_EDGE_REGION; + } + + const [region] = hostname.split('.'); + return region ? region.toLowerCase() : DEFAULT_EDGE_REGION; +} + /** * Builds a resources URL under the assumption that the API URL is in the format of https://api.axiom.co by replacing the subdomain with app. * @param urlString - The API URL @@ -27,14 +48,16 @@ const buildConsoleUrl = (urlString: string) => { export function resolveAxiomConnection( config: ResolvedAxiomConfig, consoleUrlOverride?: string, -): AxiomEvalInstrumentationOptions & { consoleEndpointUrl: string } { +): AxiomEvalInstrumentationOptions & { edgeRegion: string; consoleEndpointUrl: string } { const consoleEndpointUrl = consoleUrlOverride ?? buildConsoleUrl(config.eval.url); // Use edgeUrl for ingest/query operations, falling back to url if not specified const edgeUrl = config.eval.edgeUrl || config.eval.url; + const edgeRegion = resolveEdgeRegion(edgeUrl); return { url: config.eval.url, edgeUrl, + edgeRegion, consoleEndpointUrl: consoleEndpointUrl, token: config.eval.token, dataset: config.eval.dataset, diff --git a/packages/ai/src/config/validate-eval-token-permissions.ts b/packages/ai/src/config/validate-eval-token-permissions.ts index 42c98625..e40d5335 100644 --- a/packages/ai/src/config/validate-eval-token-permissions.ts +++ b/packages/ai/src/config/validate-eval-token-permissions.ts @@ -67,6 +67,7 @@ export async function validateTokenPermissions(config: ResolvedAxiomConfig) { try { const headers: Record = { + 'X-Axiom-Org-Id': connection.orgId ?? '', 'X-Axiom-Dataset': connection.dataset, }; @@ -74,11 +75,10 @@ export async function validateTokenPermissions(config: ResolvedAxiomConfig) { headers.Authorization = `Bearer ${connection.token}`; } - const edge = new URL(connection.edgeUrl); - const region = edge.hostname || 'us-east-1'; - console.log({ region }); + console.log({ region: connection.edgeRegion }); + const response = await fetch( - `${connection.url}/api/v3/evaluations/validate?dataset=${connection.dataset}®ion=${region}`, + `${connection.url}/api/v3/evaluations/validate?dataset=${connection.dataset}®ion=${connection.edgeRegion}`, { headers, }, @@ -88,6 +88,7 @@ export async function validateTokenPermissions(config: ResolvedAxiomConfig) { let serverMessage: string | undefined; try { const data = await response.json(); + console.debug('validation response', { data }) serverMessage = data?.error || data?.message; } catch { serverMessage = undefined; diff --git a/packages/ai/test/config/load-config.test.ts b/packages/ai/test/config/load-config.test.ts index c0479555..0f098a29 100644 --- a/packages/ai/test/config/load-config.test.ts +++ b/packages/ai/test/config/load-config.test.ts @@ -85,6 +85,7 @@ describe('resolveAxiomConnection', () => { expect(connection.edgeUrl).toBe('https://api.axiom.co'); expect(connection.url).toBe('https://api.axiom.co'); + expect(connection.edgeRegion).toBe('us-east-1'); }); it('uses edgeUrl when explicitly set', () => { @@ -96,6 +97,7 @@ describe('resolveAxiomConnection', () => { expect(connection.edgeUrl).toBe('https://eu-central-1.aws.edge.axiom.co'); expect(connection.url).toBe('https://api.axiom.co'); + expect(connection.edgeRegion).toBe('eu-central-1'); }); it('falls back to url when edgeUrl is empty string', () => { @@ -103,5 +105,33 @@ describe('resolveAxiomConnection', () => { const connection = resolveAxiomConnection(config); expect(connection.edgeUrl).toBe('https://api.axiom.co'); + expect(connection.edgeRegion).toBe('us-east-1'); + }); + + it('defaults edgeRegion to us-east-1 for non-edge hosts', () => { + const config = createConfig({ + url: 'https://api.axiom.co', + edgeUrl: 'https://api.dev.axiomtestlabs.co', + }); + const connection = resolveAxiomConnection(config); + + expect(connection.edgeRegion).toBe('us-east-1'); + }); + + it('defaults edgeRegion to us-east-1 for localhost edgeUrl', () => { + const config = createConfig({ url: 'https://api.axiom.co', edgeUrl: 'http://localhost:3000' }); + const connection = resolveAxiomConnection(config); + + expect(connection.edgeRegion).toBe('us-east-1'); + }); + + it('resolves edgeRegion when edge host does not include edge segment', () => { + const config = createConfig({ + url: 'https://api.axiom.co', + edgeUrl: 'https://us-east-1.aws.axiom.co', + }); + const connection = resolveAxiomConnection(config); + + expect(connection.edgeRegion).toBe('us-east-1'); }); });