@output.ai/core 0.4.2 → 0.4.3

package/package.json

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

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

@@ -5,16 +5,19 @@ import { EOL } from 'node:os';
 
 const oneMonthInSeconds = 60 * 60 * 24 * 30;
 
-const accumulate = async ( { entry, executionContext: { workflowName, workflowId } } ) => {
+const addEntry = async ( { entry, executionContext: { workflowName, workflowId } } ) => {
   const key = `traces/${workflowName}/${workflowId}`;
-  const transaction = ( await getRedisClient() ).multi();
-
-  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();
+  const client = await getRedisClient();
+  await client.multi()
+    .zAdd( key, [ { score: entry.timestamp, value: JSON.stringify( entry ) } ], { NX: true } )
+    .expire( key, oneMonthInSeconds, 'GT' )
+    .exec();
+};
+
+const getAllEntries = async ( { workflowName, workflowId } ) => {
+  const key = `traces/${workflowName}/${workflowId}`;
+  const client = await getRedisClient();
+  const zList = await client.zRange( key, 0, -1 );
   return zList.map( v => JSON.parse( v ) );
 };
 
@@ -41,13 +44,19 @@ export const init = async () => {
  */
 export const exec = async ( { entry, executionContext } ) => {
   const { workflowName, workflowId, startTime } = executionContext;
-  const content = buildTraceTree( await accumulate( { entry, executionContext } ) );
+
+  await addEntry( { entry, executionContext } );
 
   const isRootWorkflowEnd = !entry.parentId && entry.phase !== 'start';
-  return isRootWorkflowEnd ? upload( {
+  if ( !isRootWorkflowEnd ) {
+    return 0;
+  }
+
+  const content = buildTraceTree( await getAllEntries( { workflowName, workflowId } ) );
+  return upload( {
     key: getS3Key( { workflowId, workflowName, startTime } ),
     content: JSON.stringify( content, undefined, 2 ) + EOL
-  } ) : 0;
+  } );
 };
 
 /**

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

@@ -3,10 +3,10 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 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 getRedisClientMock = vi.fn( async () => ( { multi: () => redisMulti, zRange: zRangeMock } ) );
 vi.mock( './redis_client.js', () => ( { getRedisClient: getRedisClientMock } ) );
 
 const uploadMock = vi.fn();
@@ -32,18 +32,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,7 +49,10 @@ 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 );