save progress

Author: thesollyz

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,

packages/ai/src/config/validate-eval-token-permissions.ts

@@ -67,18 +67,18 @@ export async function validateTokenPermissions(config: ResolvedAxiomConfig) {
 
   try {
     const headers: Record<string, string> = {
+      'X-Axiom-Org-Id': connection.orgId ?? '',
       '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 });
+    console.log({ region: connection.edgeRegion });
+
     const response = await fetch(
-      `${connection.url}/api/v3/evaluations/validate?dataset=${connection.dataset}&region=${region}`,
+      `${connection.url}/api/v3/evaluations/validate?dataset=${connection.dataset}&region=${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;

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,12 +97,41 @@ 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', () => {
     const config = createConfig({ url: 'https://api.axiom.co', edgeUrl: '' });
     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');
   });
 });