@output.ai/core 0.4.2 → 0.4.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@output.ai/core",
-  "version": "0.4.2",
+  "version": "0.4.4",
   "description": "The core module of the output framework",
   "type": "module",
   "exports": {

package/src/tracing/processors/s3/configs.js

@@ -0,0 +1,29 @@
+import * as z from 'zod';
+
+const envVarSchema = z.object( {
+  OUTPUT_AWS_REGION: z.string(),
+  OUTPUT_AWS_ACCESS_KEY_ID: z.string(),
+  OUTPUT_AWS_SECRET_ACCESS_KEY: z.string(),
+  OUTPUT_TRACE_REMOTE_S3_BUCKET: z.string(),
+  OUTPUT_REDIS_URL: z.string(),
+  OUTPUT_REDIS_TRACE_TTL: z.coerce.number().int().positive().default( 60 * 60 * 24 * 7 ) // 7 days
+} );
+
+const env = {};
+
+export const loadEnv = () => {
+  const parsedFields = envVarSchema.parse( process.env );
+  env.awsRegion = parsedFields.OUTPUT_AWS_REGION;
+  env.awsAccessKeyId = parsedFields.OUTPUT_AWS_ACCESS_KEY_ID;
+  env.awsSecretAccessKey = parsedFields.OUTPUT_AWS_SECRET_ACCESS_KEY;
+  env.remoteS3Bucket = parsedFields.OUTPUT_TRACE_REMOTE_S3_BUCKET;
+  env.redisUrl = parsedFields.OUTPUT_REDIS_URL;
+  env.redisIncompleteWorkflowsTTL = parsedFields.OUTPUT_REDIS_TRACE_TTL;
+};
+
+export const getVars = () => {
+  if ( Object.keys( env ).length === 0 ) {
+    throw new Error( 'Env vars not loaded. Use loadEnv() first.' );
+  }
+  return env;
+};

package/src/tracing/processors/s3/configs.spec.js

@@ -0,0 +1,64 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+
+async function loadModule() {
+  vi.resetModules();
+  return import( './configs.js' );
+}
+
+describe( 'tracing/processors/s3/configs', () => {
+  const required = {
+    OUTPUT_AWS_REGION: 'us-east-1',
+    OUTPUT_AWS_ACCESS_KEY_ID: 'id',
+    OUTPUT_AWS_SECRET_ACCESS_KEY: 'sek',
+    OUTPUT_TRACE_REMOTE_S3_BUCKET: 'bkt',
+    OUTPUT_REDIS_URL: 'redis://localhost:6379'
+  };
+
+  beforeEach( () => {
+    vi.stubEnv( 'OUTPUT_AWS_REGION', required.OUTPUT_AWS_REGION );
+    vi.stubEnv( 'OUTPUT_AWS_ACCESS_KEY_ID', required.OUTPUT_AWS_ACCESS_KEY_ID );
+    vi.stubEnv( 'OUTPUT_AWS_SECRET_ACCESS_KEY', required.OUTPUT_AWS_SECRET_ACCESS_KEY );
+    vi.stubEnv( 'OUTPUT_TRACE_REMOTE_S3_BUCKET', required.OUTPUT_TRACE_REMOTE_S3_BUCKET );
+    vi.stubEnv( 'OUTPUT_REDIS_URL', required.OUTPUT_REDIS_URL );
+  } );
+
+  afterEach( () => {
+    vi.unstubAllEnvs();
+  } );
+
+  it( 'loadEnv() throws when required env vars are missing', async () => {
+    vi.stubEnv( 'OUTPUT_REDIS_URL', undefined );
+    const { loadEnv } = await loadModule();
+    expect( () => loadEnv() ).toThrow( /OUTPUT_REDIS_URL/ );
+  } );
+
+  it( 'loadEnv() populates getVars() with parsed env', async () => {
+    const { loadEnv, getVars } = await loadModule();
+    loadEnv();
+    const vars = getVars();
+    expect( vars.awsRegion ).toBe( required.OUTPUT_AWS_REGION );
+    expect( vars.awsAccessKeyId ).toBe( required.OUTPUT_AWS_ACCESS_KEY_ID );
+    expect( vars.awsSecretAccessKey ).toBe( required.OUTPUT_AWS_SECRET_ACCESS_KEY );
+    expect( vars.remoteS3Bucket ).toBe( required.OUTPUT_TRACE_REMOTE_S3_BUCKET );
+    expect( vars.redisUrl ).toBe( required.OUTPUT_REDIS_URL );
+    expect( vars.redisIncompleteWorkflowsTTL ).toBe( 60 * 60 * 24 * 7 );
+  } );
+
+  it( 'loadEnv() uses OUTPUT_REDIS_TRACE_TTL when set', async () => {
+    vi.stubEnv( 'OUTPUT_REDIS_TRACE_TTL', '3600' );
+    const { loadEnv, getVars } = await loadModule();
+    loadEnv();
+    expect( getVars().redisIncompleteWorkflowsTTL ).toBe( 3600 );
+  } );
+
+  it( 'getVars() throws when loadEnv() was not called', async () => {
+    const { getVars } = await loadModule();
+    expect( () => getVars() ).toThrow( 'Env vars not loaded. Use loadEnv() first.' );
+  } );
+
+  it( 'loadEnv() throws when OUTPUT_REDIS_TRACE_TTL is invalid', async () => {
+    vi.stubEnv( 'OUTPUT_REDIS_TRACE_TTL', 'not-a-number' );
+    const { loadEnv } = await loadModule();
+    expect( () => loadEnv() ).toThrow();
+  } );
+} );

package/src/tracing/processors/s3/index.js

@@ -2,22 +2,51 @@ import { upload } from './s3_client.js';
 import { getRedisClient } from './redis_client.js';
 import buildTraceTree from '../../tools/build_trace_tree.js';
 import { EOL } from 'node:os';
+import { loadEnv, getVars } from './configs.js';
 
-const oneMonthInSeconds = 60 * 60 * 24 * 30;
+const createRedisKey = ( { workflowId, workflowName } ) => `traces/${workflowName}/${workflowId}`;
 
-const accumulate = async ( { entry, executionContext: { workflowName, workflowId } } ) => {
-  const key = `traces/${workflowName}/${workflowId}`;
-  const transaction = ( await getRedisClient() ).multi();
+/**
+ * Add new entry to list of entries
+ * @param {object} entry
+ * @param {string} key
+ */
+const addEntry = async ( entry, key ) => {
+  const client = await getRedisClient();
+  await client.multi()
+    .zAdd( key, [ { score: entry.timestamp, value: JSON.stringify( entry ) } ], { NX: true } )
+    .expire( key, getVars().redisIncompleteWorkflowsTTL, 'GT' )
+    .exec();
+};
 
-  transaction.zAdd( key, [
-    { score: entry.timestamp, value: JSON.stringify( entry ) }
-  ], { NX: true } );
-  transaction.expire( key, oneMonthInSeconds, 'GT' );
-  transaction.zRange( key, 0, -1 );
-  const [ ,, zList ] = await transaction.exec();
+/**
+ * Returns entries from cache, parsed to object
+ * @param {string} key
+ * @returns {object[]}
+ */
+const getEntries = async key => {
+  const client = await getRedisClient();
+  const zList = await client.zRange( key, 0, -1 );
   return zList.map( v => JSON.parse( v ) );
 };
 
+/**
+ * Removes the entries from cache
+ * @param {string} key
+ */
+const bustEntries = async key => {
+  const client = await getRedisClient();
+  await client.del( key );
+};
+
+/**
+ * Return the S3 key for the trace file
+ * @param {object} args
+ * @param {number} args.startTime
+ * @param {string} args.workflowId
+ * @param {string} args.workflowName
+ * @returns
+ */
 const getS3Key = ( { startTime, workflowId, workflowName } ) => {
   const isoDate = new Date( startTime ).toISOString();
   const [ year, month, day ] = isoDate.split( /\D/, 3 );
@@ -29,6 +58,7 @@ const getS3Key = ( { startTime, workflowId, workflowName } ) => {
  * Init this processor
  */
 export const init = async () => {
+  loadEnv();
   await getRedisClient();
 };
 
@@ -41,13 +71,21 @@ export const init = async () => {
  */
 export const exec = async ( { entry, executionContext } ) => {
   const { workflowName, workflowId, startTime } = executionContext;
-  const content = buildTraceTree( await accumulate( { entry, executionContext } ) );
+  const cacheKey = createRedisKey( { workflowId, workflowName } );
+
+  await addEntry( entry, cacheKey );
 
   const isRootWorkflowEnd = !entry.parentId && entry.phase !== 'start';
-  return isRootWorkflowEnd ? upload( {
+  if ( !isRootWorkflowEnd ) {
+    return;
+  }
+
+  const content = buildTraceTree( await getEntries( cacheKey ) );
+  await upload( {
     key: getS3Key( { workflowId, workflowName, startTime } ),
     content: JSON.stringify( content, undefined, 2 ) + EOL
-  } ) : 0;
+  } );
+  await bustEntries( cacheKey );
 };
 
 /**
@@ -59,4 +97,4 @@ export const exec = async ( { entry, executionContext } ) => {
  * @returns {string} The S3 url of the trace file
  */
 export const getDestination = ( { startTime, workflowId, workflowName } ) =>
-  `https://${process.env.OUTPUT_TRACE_REMOTE_S3_BUCKET}.s3.amazonaws.com/${getS3Key( { workflowId, workflowName, startTime } )}`;
+  `https://${getVars().remoteS3Bucket}.s3.amazonaws.com/${getS3Key( { workflowId, workflowName, startTime } )}`;

package/src/tracing/processors/s3/index.spec.js

@@ -1,12 +1,24 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 
+const loadEnvMock = vi.fn();
+const getVarsMock = vi.fn( () => ( {
+  remoteS3Bucket: 'bkt',
+  redisIncompleteWorkflowsTTL: 3600
+} ) );
+vi.mock( './configs.js', () => ( { loadEnv: loadEnvMock, getVars: getVarsMock } ) );
+
 const redisMulti = {
   zAdd: vi.fn().mockReturnThis(),
   expire: vi.fn().mockReturnThis(),
-  zRange: vi.fn().mockReturnThis(),
   exec: vi.fn()
 };
-const getRedisClientMock = vi.fn( async () => ( { multi: () => redisMulti } ) );
+const zRangeMock = vi.fn();
+const delMock = vi.fn().mockResolvedValue( undefined );
+const getRedisClientMock = vi.fn( async () => ( {
+  multi: () => redisMulti,
+  zRange: zRangeMock,
+  del: delMock
+} ) );
 vi.mock( './redis_client.js', () => ( { getRedisClient: getRedisClientMock } ) );
 
 const uploadMock = vi.fn();
@@ -18,12 +30,13 @@ vi.mock( '../../tools/build_trace_tree.js', () => ( { default: buildTraceTreeMoc
 describe( 'tracing/processors/s3', () => {
   beforeEach( () => {
     vi.clearAllMocks();
-    process.env.OUTPUT_TRACE_REMOTE_S3_BUCKET = 'bkt';
+    getVarsMock.mockReturnValue( { remoteS3Bucket: 'bkt', redisIncompleteWorkflowsTTL: 3600 } );
   } );
 
-  it( 'init(): ensures redis client is created', async () => {
+  it( 'init(): loads config and ensures redis client is created', async () => {
     const { init } = await import( './index.js' );
     await init();
+    expect( loadEnvMock ).toHaveBeenCalledTimes( 1 );
     expect( getRedisClientMock ).toHaveBeenCalledTimes( 1 );
   } );
 
@@ -32,18 +45,15 @@ describe( 'tracing/processors/s3', () => {
     const startTime = Date.parse( '2020-01-02T03:04:05.678Z' );
     const ctx = { executionContext: { workflowId: 'id1', workflowName: 'WF', startTime } };
 
-    // Redis will return progressively larger sorted sets
-    redisMulti.exec
-      .mockResolvedValueOnce( [ , , [ JSON.stringify( { name: 'A', phase: 'start', timestamp: startTime } ) ] ] )
-      .mockResolvedValueOnce( [ , , [
-        JSON.stringify( { name: 'A', phase: 'start', timestamp: startTime } ),
-        JSON.stringify( { name: 'A', phase: 'tick', timestamp: startTime + 1 } )
-      ] ] )
-      .mockResolvedValueOnce( [ , , [
-        JSON.stringify( { name: 'A', phase: 'start', timestamp: startTime } ),
-        JSON.stringify( { name: 'A', phase: 'tick', timestamp: startTime + 1 } ),
-        JSON.stringify( { name: 'A', phase: 'end', timestamp: startTime + 2 } )
-      ] ] );
+    // multi().exec() just needs to resolve for addEntry calls
+    redisMulti.exec.mockResolvedValue( [] );
+
+    // zRange is only called once at the end to get all entries
+    zRangeMock.mockResolvedValue( [
+      JSON.stringify( { name: 'A', phase: 'start', timestamp: startTime } ),
+      JSON.stringify( { name: 'A', phase: 'tick', timestamp: startTime + 1 } ),
+      JSON.stringify( { name: 'A', phase: 'end', timestamp: startTime + 2 } )
+    ] );
 
     await exec( { ...ctx, entry: { name: 'A', phase: 'start', timestamp: startTime, parentId: 'root' } } );
     await exec( { ...ctx, entry: { name: 'A', phase: 'tick', timestamp: startTime + 1, parentId: 'root' } } );
@@ -52,13 +62,30 @@ describe( 'tracing/processors/s3', () => {
 
     // Accumulation happened 3 times
     expect( redisMulti.zAdd ).toHaveBeenCalledTimes( 3 );
-    expect( buildTraceTreeMock ).toHaveBeenCalledTimes( 3 );
+
+    // Tree is only built once at the end (not on every event)
+    expect( buildTraceTreeMock ).toHaveBeenCalledTimes( 1 );
+    expect( zRangeMock ).toHaveBeenCalledTimes( 1 );
 
     // Only last call triggers upload
     expect( uploadMock ).toHaveBeenCalledTimes( 1 );
     const { key, content } = uploadMock.mock.calls[0][0];
     expect( key ).toMatch( /^WF\/2020\/01\/02\// );
     expect( JSON.parse( content.trim() ).count ).toBe( 3 );
+
+    expect( delMock ).toHaveBeenCalledTimes( 1 );
+    expect( delMock ).toHaveBeenCalledWith( 'traces/WF/id1' );
+  } );
+
+  it( 'getDestination(): returns S3 URL using bucket and key from getVars', async () => {
+    getVarsMock.mockReturnValue( { remoteS3Bucket: 'my-bucket', redisIncompleteWorkflowsTTL: 3600 } );
+    const { getDestination } = await import( './index.js' );
+    const startTime = Date.parse( '2020-01-02T03:04:05.678Z' );
+    const url = getDestination( { workflowId: 'id1', workflowName: 'WF', startTime } );
+    expect( getVarsMock ).toHaveBeenCalled();
+    expect( url ).toBe(
+      'https://my-bucket.s3.amazonaws.com/WF/2020/01/02/2020-01-02-03-04-05-678Z_id1.json'
+    );
   } );
 } );
 

package/src/tracing/processors/s3/redis_client.js

@@ -1,6 +1,6 @@
 import { createClient } from 'redis';
-import { throws } from '#utils';
 import { createChildLogger } from '#logger';
+import { getVars } from './configs.js';
 
 const log = createChildLogger( 'RedisClient' );
 
@@ -36,11 +36,10 @@ async function connect( url ) {
  * Concurrent calls during connection will receive the same pending promise.
  *
  * @returns {Promise<redis.RedisClientType>} Connected Redis client
- * @throws {Error} If OUTPUT_REDIS_URL env var is missing
  * @throws {Error} If connection fails (wrapped with context)
  */
 export async function getRedisClient() {
-  const url = process.env.OUTPUT_REDIS_URL ?? throws( new Error( 'Missing OUTPUT_REDIS_URL environment variable' ) );
+  const url = getVars().redisUrl;
 
   const pingResult = await state.client?.ping().catch( err => {
     log.error( 'Redis ping failed', { error: err.message, code: err.code } );

package/src/tracing/processors/s3/redis_client.spec.js

@@ -14,6 +14,9 @@ vi.mock( '#logger', () => ( {
   } )
 } ) );
 
+const getVarsMock = vi.fn();
+vi.mock( './configs.js', () => ( { getVars: () => getVarsMock() } ) );
+
 const createClientImpl = vi.fn();
 vi.mock( 'redis', () => ( { createClient: opts => createClientImpl( opts ) } ) );
 
@@ -25,7 +28,7 @@ async function loadModule() {
 describe( 'tracing/processors/s3/redis_client', () => {
   beforeEach( () => {
     vi.clearAllMocks();
-    delete process.env.OUTPUT_REDIS_URL;
+    getVarsMock.mockReturnValue( {} );
     logCalls.warn = [];
     logCalls.error = [];
   } );
@@ -34,13 +37,14 @@ describe( 'tracing/processors/s3/redis_client', () => {
     vi.useRealTimers();
   } );
 
-  it( 'throws if OUTPUT_REDIS_URL is missing', async () => {
+  it( 'throws when config redisUrl is missing', async () => {
+    getVarsMock.mockReturnValue( {} );
     const { getRedisClient } = await loadModule();
-    await expect( getRedisClient() ).rejects.toThrow( 'Missing OUTPUT_REDIS_URL' );
+    await expect( getRedisClient() ).rejects.toThrow();
   } );
 
   it( 'creates client with url, connects once, then reuses cached when ping is PONG', async () => {
-    process.env.OUTPUT_REDIS_URL = 'redis://localhost:6379';
+    getVarsMock.mockReturnValue( { redisUrl: 'redis://localhost:6379' } );
 
     const pingMock = vi.fn().mockResolvedValue( 'PONG' );
     const connectMock = vi.fn().mockResolvedValue();
@@ -63,7 +67,7 @@ describe( 'tracing/processors/s3/redis_client', () => {
   } );
 
   it( 'closes stale client and reconnects when ping fails', async () => {
-    process.env.OUTPUT_REDIS_URL = 'redis://localhost:6379';
+    getVarsMock.mockReturnValue( { redisUrl: 'redis://localhost:6379' } );
 
     const quitMock = vi.fn().mockResolvedValue();
     const connectMock = vi.fn().mockResolvedValue();
@@ -92,7 +96,7 @@ describe( 'tracing/processors/s3/redis_client', () => {
   } );
 
   it( 'reconnects successfully even when quit() on stale client rejects', async () => {
-    process.env.OUTPUT_REDIS_URL = 'redis://localhost:6379';
+    getVarsMock.mockReturnValue( { redisUrl: 'redis://localhost:6379' } );
 
     const quitMock = vi.fn().mockRejectedValue( new Error( 'Quit failed' ) );
     const connectMock = vi.fn().mockResolvedValue();
@@ -121,7 +125,7 @@ describe( 'tracing/processors/s3/redis_client', () => {
   } );
 
   it( 'wraps connect() errors with code and cleans up failed client', async () => {
-    process.env.OUTPUT_REDIS_URL = 'redis://localhost:6379';
+    getVarsMock.mockReturnValue( { redisUrl: 'redis://localhost:6379' } );
 
     const connectErr = new Error( 'Connection refused' );
     connectErr.code = 'ECONNREFUSED';
@@ -147,7 +151,7 @@ describe( 'tracing/processors/s3/redis_client', () => {
   } );
 
   it( 'logs ping failures with error level', async () => {
-    process.env.OUTPUT_REDIS_URL = 'redis://localhost:6379';
+    getVarsMock.mockReturnValue( { redisUrl: 'redis://localhost:6379' } );
 
     const pingErr = new Error( 'Connection reset' );
     pingErr.code = 'ECONNRESET';

package/src/tracing/processors/s3/s3_client.js

@@ -1,5 +1,5 @@
 import { S3Client, PutObjectCommand } from '@aws-sdk/client-s3';
-import { throws } from '#utils';
+import { getVars } from './configs.js';
 
 const state = { s3Client: null };
 
@@ -12,9 +12,7 @@ const getS3Client = () => {
     return state.s3Client;
   }
 
-  const region = process.env.OUTPUT_AWS_REGION ?? throws( new Error( 'Missing OUTPUT_AWS_REGION env var' ) );
-  const secretAccessKey = process.env.OUTPUT_AWS_SECRET_ACCESS_KEY ?? throws( new Error( 'Missing OUTPUT_AWS_SECRET_ACCESS_KEY env var' ) );
-  const accessKeyId = process.env.OUTPUT_AWS_ACCESS_KEY_ID ?? throws( new Error( 'Missing OUTPUT_AWS_ACCESS_KEY_ID env var' ) );
+  const { awsRegion: region, awsSecretAccessKey: secretAccessKey, awsAccessKeyId: accessKeyId } = getVars();
 
   return state.s3Client = new S3Client( { region, credentials: { accessKeyId, secretAccessKey } } );
 };
@@ -26,8 +24,4 @@ const getS3Client = () => {
  * @param {string} content - File content
  */
 export const upload = ( { key, content } ) =>
-  getS3Client().send( new PutObjectCommand( {
-    Bucket: process.env.OUTPUT_TRACE_REMOTE_S3_BUCKET ?? throws( new Error( 'Missing OUTPUT_TRACE_REMOTE_S3_BUCKET env var' ) ),
-    Key: key,
-    Body: content
-  } ) );
+  getS3Client().send( new PutObjectCommand( { Bucket: getVars().remoteS3Bucket, Key: key, Body: content } ) );

package/src/tracing/processors/s3/s3_client.spec.js

@@ -6,6 +6,9 @@ vi.mock( '#utils', () => ( {
   }
 } ) );
 
+const getVarsMock = vi.fn();
+vi.mock( './configs', () => ( { getVars: () => getVarsMock() } ) );
+
 const sendMock = vi.fn();
 const ctorState = { args: null };
 class S3ClientMock {
@@ -32,23 +35,15 @@ async function loadModule() {
 describe( 'tracing/processors/s3/s3_client', () => {
   beforeEach( () => {
     vi.clearAllMocks();
-    delete process.env.OUTPUT_AWS_REGION;
-    delete process.env.OUTPUT_AWS_SECRET_ACCESS_KEY;
-    delete process.env.OUTPUT_AWS_ACCESS_KEY_ID;
-    delete process.env.OUTPUT_TRACE_REMOTE_S3_BUCKET;
+    getVarsMock.mockReturnValue( {
+      awsRegion: 'us-east-1',
+      awsAccessKeyId: 'id',
+      awsSecretAccessKey: 'sek',
+      remoteS3Bucket: 'bucket'
+    } );
   } );
 
-  it( 'fails fast when required env vars are missing for client creation', async () => {
-    const { upload } = await loadModule();
-    expect( () => upload( { key: 'k', content: 'c' } ) ).toThrow();
-  } );
-
-  it( 'creates client once with env and uploads with bucket/key/content', async () => {
-    process.env.OUTPUT_AWS_REGION = 'us-east-1';
-    process.env.OUTPUT_AWS_SECRET_ACCESS_KEY = 'sek';
-    process.env.OUTPUT_AWS_ACCESS_KEY_ID = 'id';
-    process.env.OUTPUT_TRACE_REMOTE_S3_BUCKET = 'bucket';
-
+  it( 'creates client once with config and uploads with bucket/key/content', async () => {
     const { upload } = await loadModule();
 
     await upload( { key: 'wf/key.json', content: '{"a":1}' } );