@outputai/core 0.5.1-dev.45fb889.0 → 0.5.1-next.8e45051.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/core",
-  "version": "0.5.1-dev.45fb889.0",
+  "version": "0.5.1-next.8e45051.0",
   "description": "The core module of the output framework",
   "type": "module",
   "exports": {

package/src/tracing/trace_engine.js

@@ -4,7 +4,7 @@ import { serializeError } from './tools/utils.js';
 import { isStringboolTrue } from '#utils';
 import * as localProcessor from './processors/local/index.js';
 import * as s3Processor from './processors/s3/index.js';
-import { ComponentType, Signal } from '#consts';
+import { ComponentType } from '#consts';
 import { createChildLogger } from '#logger';
 import { EventAction } from './trace_consts.js';
 import { BaseAttribute } from './trace_attribute.js';
@@ -93,14 +93,14 @@ export const addEventAction = ( action, { kind, name, id, parentId, details, exe
 export function addEventActionWithContext( action, options ) {
   const storeContent = Storage.load();
   if ( storeContent ) { // If there is no storageContext this was not called from a Temporal environment
-    const { parentId, parentName, executionContext, workflowHandle } = storeContent;
+    const { parentId, executionContext, sendAttributeSignal } = storeContent;
     if ( action === EventAction.ADD_ATTR ) {
       const attribute = options.details;
       if ( !( attribute instanceof BaseAttribute ) ) {
-        throw new Error( `${EventAction.ADD_ATTR} called argument that is not a BaseAttribute instance` );
+        throw new Error( `Event ${EventAction.ADD_ATTR} argument is not a BaseAttribute instance` );
+      } else {
+        sendAttributeSignal( options.details );
       }
-      attribute.setActivity( parentId, parentName );
-      workflowHandle.signal( Signal.ADD_ATTRIBUTE, attribute );
     }
     addEventAction( action, { ...options, parentId, executionContext } );
   }

package/src/tracing/trace_engine.spec.js

@@ -5,6 +5,12 @@ vi.mock( '#async_storage', () => ( {
   Storage: { load: storageLoadMock }
 } ) );
 
+const logWarnMock = vi.fn();
+const logErrorMock = vi.fn();
+vi.mock( '#logger', () => ( {
+  createChildLogger: () => ( { warn: logWarnMock, error: logErrorMock } )
+} ) );
+
 const localInitMock = vi.fn( async () => {} );
 const localExecMock = vi.fn();
 const localGetDestinationMock = vi.fn( () => '/local/path.json' );
@@ -117,6 +123,62 @@ describe( 'tracing/trace_engine', () => {
     expect( payload.entry.action ).toBe( 'tick' );
   } );
 
+  it( 'addEventActionWithContext() sends ADD_ATTR attributes through storage context', async () => {
+    process.env.OUTPUT_TRACE_LOCAL_ON = 'true';
+    const sendAttributeSignalMock = vi.fn();
+    const executionContext = { runId: 'r1', disableTrace: false };
+    storageLoadMock.mockReturnValue( {
+      parentId: 'ctx-p',
+      executionContext,
+      sendAttributeSignal: sendAttributeSignalMock
+    } );
+    const { init, addEventActionWithContext } = await loadTraceEngine();
+    const { EventAction } = await import( './trace_consts.js' );
+    const { Attribute } = await import( './trace_attribute.js' );
+    await init();
+
+    const attribute = new Attribute.HTTPRequestCount( 'https://example.test', 'req-1' );
+    addEventActionWithContext( EventAction.ADD_ATTR, { kind: 'http', name: 'request', id: 'req-1', details: attribute } );
+
+    expect( sendAttributeSignalMock ).toHaveBeenCalledTimes( 1 );
+    expect( sendAttributeSignalMock ).toHaveBeenCalledWith( attribute );
+    expect( localExecMock ).toHaveBeenCalledTimes( 1 );
+    expect( localExecMock.mock.calls[0][0] ).toEqual( {
+      executionContext,
+      entry: {
+        kind: 'http',
+        action: EventAction.ADD_ATTR,
+        name: 'request',
+        id: 'req-1',
+        parentId: 'ctx-p',
+        timestamp: expect.any( Number ),
+        details: attribute
+      }
+    } );
+  } );
+
+  it( 'addEventActionWithContext() throws on invalid ADD_ATTR signal payloads', async () => {
+    process.env.OUTPUT_TRACE_LOCAL_ON = 'true';
+    const sendAttributeSignalMock = vi.fn();
+    storageLoadMock.mockReturnValue( {
+      parentId: 'ctx-p',
+      executionContext: { runId: 'r1', disableTrace: false },
+      sendAttributeSignal: sendAttributeSignalMock
+    } );
+    const { init, addEventActionWithContext } = await loadTraceEngine();
+    const { EventAction } = await import( './trace_consts.js' );
+    await init();
+
+    const invalidAttribute = { type: 'not-a-base-attribute' };
+    expect( () => addEventActionWithContext(
+      EventAction.ADD_ATTR,
+      { kind: 'http', name: 'request', id: 'req-1', details: invalidAttribute }
+    ) ).toThrow( /not a BaseAttribute instance/ );
+
+    expect( sendAttributeSignalMock ).not.toHaveBeenCalled();
+    expect( localExecMock ).not.toHaveBeenCalled();
+  } );
+
   it( 'addEventActionWithContext() does not emit when storage executionContext.disableTrace is true', async () => {
     process.env.OUTPUT_TRACE_LOCAL_ON = '1';
     storageLoadMock.mockReturnValue( {

package/src/utils/index.d.ts

@@ -132,3 +132,16 @@ export function deepMerge( a: object, b: object ): object;
  * @returns Short string using A–Z, a–z, 0–9, `_`, `-` (typically 21–22 chars).
  */
 export function toUrlSafeBase64( uuid: string ): string;
+
+/**
+ * Similar to native Promise.allSettled, but rejects with `{ isTimeout: true }`
+ * if the execution exceeds the given timeout.
+ *
+ * @param promises - Values or promises to wait for.
+ * @param timeoutMs - Maximum wait time in milliseconds.
+ * @returns Native Promise.allSettled-style results.
+ */
+export function allSettledWithTimeout<T>(
+  promises: Array<T | PromiseLike<T>>,
+  timeoutMs: number
+): Promise<PromiseSettledResult<Awaited<T>>[]>;

package/src/utils/utils.js

@@ -209,3 +209,42 @@ export const toUrlSafeBase64 = uuid => {
   const toDigits = n => n <= 0n ? [] : toDigits( n / base ).concat( alphabet[Number( n % base )] );
   return toDigits( BigInt( '0x' + hex ) ).join( '' );
 };
+
+/**
+ * Similar to native Promise.allSettled but throws an Error if the execution exceeds a given time.
+ *
+ * The error thrown will have attribute `.isTimeout` as `true`.
+ *
+ * @template T
+ * @param {Array<T | PromiseLike<T>>} promises
+ * @param {number} timeoutMs
+ * @returns {Promise<PromiseSettledResult<Awaited<T>>[]>}
+ * @throws {Error & { isTimeout: true }}
+ */
+export const allSettledWithTimeout = ( () => {
+  class TimeoutError extends Error {
+    isTimeout = true;
+    constructor() {
+      super( 'Timed out before completing all promises' );
+    }
+  }
+
+  return async ( promises, timeoutMs ) => {
+    if ( promises.length === 0 ) {
+      return [];
+    }
+
+    const state = { timeoutMonitor: null };
+
+    try {
+      return await Promise.race( [
+        Promise.allSettled( promises ),
+        new Promise( ( _, reject ) => {
+          state.timeoutMonitor = setTimeout( () => reject( new TimeoutError() ), timeoutMs );
+        } )
+      ] );
+    } finally {
+      clearTimeout( state.timeoutMonitor );
+    }
+  };
+} )();

package/src/utils/utils.spec.js

@@ -1,4 +1,4 @@
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, vi } from 'vitest';
 import { Readable } from 'node:stream';
 import {
   clone,
@@ -6,7 +6,8 @@ import {
   serializeFetchResponse,
   deepMerge,
   isPlainObject,
-  toUrlSafeBase64
+  toUrlSafeBase64,
+  allSettledWithTimeout
 } from './utils.js';
 
 describe( 'clone', () => {
@@ -22,6 +23,44 @@ describe( 'clone', () => {
   } );
 } );
 
+describe( 'allSettledWithTimeout', () => {
+  it( 'returns an empty array when no promises are provided', async () => {
+    await expect( allSettledWithTimeout( [], 100 ) ).resolves.toEqual( [] );
+  } );
+
+  it( 'returns native allSettled results when all promises settle before timeout', async () => {
+    const error = new Error( 'boom' );
+    const result = await allSettledWithTimeout( [
+      Promise.resolve( 'ok' ),
+      Promise.reject( error ),
+      42
+    ], 100 );
+
+    expect( result ).toEqual( [
+      { status: 'fulfilled', value: 'ok' },
+      { status: 'rejected', reason: error },
+      { status: 'fulfilled', value: 42 }
+    ] );
+  } );
+
+  it( 'rejects with a timeout-shaped error when promises do not settle in time', async () => {
+    vi.useFakeTimers();
+    try {
+      const pending = new Promise( () => {} );
+      const result = allSettledWithTimeout( [ pending ], 1000 );
+      const assertion = expect( result ).rejects.toMatchObject( {
+        isTimeout: true,
+        message: 'Timed out before completing all promises'
+      } );
+
+      await vi.advanceTimersByTimeAsync( 1000 );
+      await assertion;
+    } finally {
+      vi.useRealTimers();
+    }
+  } );
+} );
+
 describe( 'serializeFetchResponse', () => {
   it( 'serializes JSON response body and flattens headers', async () => {
     const payload = { a: 1, b: 'two' };

package/src/worker/interceptors/activity.js

@@ -2,10 +2,16 @@ import { Context } from '@temporalio/activity';
 import { Storage } from '#async_storage';
 import * as Tracing from '#tracing';
 import { headersToObject } from '../sandboxed_utils.js';
-import { BusEventType, METADATA_ACCESS_SYMBOL } from '#consts';
-import { activityHeartbeatEnabled, activityHeartbeatIntervalMs } from '../configs.js';
+import { BusEventType, METADATA_ACCESS_SYMBOL, Signal } from '#consts';
+import { activityHeartbeatEnabled, activityHeartbeatIntervalMs, namespace } from '../configs.js';
 import { messageBus } from '#bus';
 import { Client } from '@temporalio/client';
+import { createChildLogger } from '#logger';
+import { allSettledWithTimeout } from '#utils';
+
+const log = createChildLogger( 'ActivityInterceptor' );
+
+const IN_FLIGHT_SIGNALS_TIMEOUT_MS = 30_000;
 
 /*
   This interceptor wraps every activity execution with cross-cutting concerns:
@@ -36,38 +42,87 @@ export class ActivityExecutionInterceptor {
     this.connection = connection;
   };
 
+  /**
+   * Returns a workflow entry by its name or throws error
+   * @param {string} workflowName
+   * @returns {object} Workflow entry
+   * @throws {Error}
+   */
+  getWorkflowEntry( workflowName ) {
+    const workflowEntry = this.workflowsMap.get( workflowName );
+    if ( !workflowEntry ) {
+      throw new Error( `Activity interceptor: workflow "${workflowName}" not found in workflowsMap.` );
+    }
+    return workflowEntry;
+  }
+
   async execute( input, next ) {
     const startDate = Date.now();
-    const client = new Client( { connection: this.connection } );
+    const client = new Client( { connection: this.connection, namespace } );
 
     const { workflowExecution: { workflowId }, activityId: id, activityType: name, workflowType: workflowName } = Context.current().info;
     const { executionContext } = headersToObject( input.headers );
     const { type: kind } = this.activities?.[name]?.[METADATA_ACCESS_SYMBOL];
+    const { path: workflowFilename } = this.getWorkflowEntry( workflowName );
 
     const workflowHandle = client.workflow.getHandle( workflowId );
 
+    const state = {
+      heartbeat: null,
+      activityOutput: undefined,
+      signals: []
+    };
+
+    const errorContext = {
+      workflowId,
+      workflowName,
+      activityId: id,
+      activityName: name
+    };
+
+    const sendAttributeSignal = attribute => {
+      attribute.setActivity( id, name );
+      state.signals.push(
+        workflowHandle
+          .signal( Signal.ADD_ATTRIBUTE, attribute )
+          .catch( e =>
+            log.warn( `Signal "${Signal.ADD_ATTRIBUTE}" failed`, { message: e.message, stack: e.stack, activityId: id, ...errorContext } )
+          )
+      );
+    };
+
+    const flushSignals = async signals => {
+      try {
+        await allSettledWithTimeout( signals, IN_FLIGHT_SIGNALS_TIMEOUT_MS );
+      } catch ( error ) {
+        if ( error.isTimeout ) {
+          log.warn( 'Some usage/cost attributes were missed because not all activity signals were sent to the workflow', errorContext );
+        } else {
+          throw error;
+        }
+      }
+    };
+
+    // Wraps the execution with accessible metadata for the activity
+    const ctx = { parentId: id, executionContext, workflowFilename, sendAttributeSignal };
+
     messageBus.emit( BusEventType.ACTIVITY_START, { id, name, kind, workflowId, workflowName } );
     Tracing.addEventStart( { id, name, kind, parentId: workflowId, details: input.args[0], executionContext } );
 
-    const workflowEntry = this.workflowsMap.get( workflowName );
-    if ( !workflowEntry ) {
-      const availableWorkflows = [ ...this.workflowsMap.keys() ].join( ', ' );
-      throw new Error( `Activity interceptor: workflow "${workflowName}" not found in workflowsMap. Available: [${availableWorkflows}]` );
-    }
-    const workflowFilename = workflowEntry.path;
-
-    const intervals = { heartbeat: null };
     try {
       // Sends heartbeat to communicate that activity is still alive
-      intervals.heartbeat = activityHeartbeatEnabled && setInterval( () => Context.current().heartbeat(), activityHeartbeatIntervalMs );
+      state.heartbeat = activityHeartbeatEnabled && setInterval( () => Context.current().heartbeat(), activityHeartbeatIntervalMs );
 
-      // Wraps the execution with accessible metadata for the activity
-      const ctx = { parentId: id, parentName: name, executionContext, workflowFilename, workflowHandle };
-      const output = await Storage.runWithContext( async _ => next( input ), ctx );
+      try {
+        state.activityOutput = await Storage.runWithContext( async _ => next( input ), ctx );
+      } finally {
+        // Ensure in-flight signals are delivered (up to a reasonable time) before handling errors
+        await flushSignals( state.signals );
+      }
 
       messageBus.emit( BusEventType.ACTIVITY_END, { id, name, kind, workflowId, workflowName, duration: Date.now() - startDate } );
-      Tracing.addEventEnd( { id, details: output, executionContext } );
-      return output;
+      Tracing.addEventEnd( { id, details: state.activityOutput, executionContext } );
+      return state.activityOutput;
 
     } catch ( error ) {
       messageBus.emit( BusEventType.ACTIVITY_ERROR, { id, name, kind, workflowId, workflowName, duration: Date.now() - startDate, error } );
@@ -75,7 +130,7 @@ export class ActivityExecutionInterceptor {
 
       throw error;
     } finally {
-      clearInterval( intervals.heartbeat );
+      clearInterval( state.heartbeat );
     }
   }
 };

package/src/worker/interceptors/activity.spec.js

@@ -1,9 +1,12 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { BusEventType } from '#consts';
+import { BusEventType, Signal } from '#consts';
 
 const METADATA_ACCESS_SYMBOL = vi.hoisted( () => Symbol( '__metadata' ) );
 const workflowHandleMock = vi.hoisted( () => ( { signal: vi.fn() } ) );
 const getHandleMock = vi.hoisted( () => vi.fn( () => workflowHandleMock ) );
+const clientConstructorMock = vi.hoisted( () => vi.fn() );
+const allSettledWithTimeoutMock = vi.hoisted( () => vi.fn().mockResolvedValue( [] ) );
+const logWarnMock = vi.hoisted( () => vi.fn() );
 
 const heartbeatMock = vi.fn();
 const runWithContextMock = vi.hoisted( () => vi.fn().mockImplementation( async fn => fn() ) );
@@ -25,6 +28,10 @@ vi.mock( '@temporalio/activity', () => ( {
 
 vi.mock( '@temporalio/client', () => ( {
   Client: class Client {
+    constructor( options ) {
+      clientConstructorMock( options );
+    }
+
     workflow = {
       getHandle: getHandleMock
     };
@@ -37,6 +44,15 @@ vi.mock( '#async_storage', () => ( {
   }
 } ) );
 
+vi.mock( '#utils', async importOriginal => {
+  const actual = await importOriginal();
+  return { ...actual, allSettledWithTimeout: allSettledWithTimeoutMock };
+} );
+
+vi.mock( '#logger', () => ( {
+  createChildLogger: () => ( { warn: logWarnMock } )
+} ) );
+
 const addEventStartMock = vi.fn();
 const addEventEndMock = vi.fn();
 const addEventErrorMock = vi.fn();
@@ -68,6 +84,9 @@ vi.mock( '../configs.js', () => ( {
   },
   get activityHeartbeatIntervalMs() {
     return parseInt( process.env.OUTPUT_ACTIVITY_HEARTBEAT_INTERVAL_MS || '120000', 10 );
+  },
+  get namespace() {
+    return process.env.TEMPORAL_NAMESPACE || 'default';
   }
 } ) );
 
@@ -85,6 +104,8 @@ const makeInput = () => ( {
 describe( 'ActivityExecutionInterceptor', () => {
   beforeEach( () => {
     vi.clearAllMocks();
+    allSettledWithTimeoutMock.mockResolvedValue( [] );
+    workflowHandleMock.signal.mockResolvedValue( undefined );
     vi.useFakeTimers();
     vi.resetModules();
     // Default: heartbeat enabled with 50ms interval for fast tests
@@ -116,19 +137,88 @@ describe( 'ActivityExecutionInterceptor', () => {
     expect( addEventStartMock ).toHaveBeenCalledOnce();
     expect( addEventEndMock ).toHaveBeenCalledOnce();
     expect( addEventErrorMock ).not.toHaveBeenCalled();
+    expect( clientConstructorMock ).toHaveBeenCalledWith( { connection: undefined, namespace: 'default' } );
     expect( runWithContextMock ).toHaveBeenCalledWith(
       expect.any( Function ),
       expect.objectContaining( {
         parentId: 'act-1',
-        parentName: 'myWorkflow#myStep',
         executionContext: { workflowId: 'wf-1' },
         workflowFilename: '/workflows/myWorkflow.js',
-        workflowHandle: workflowHandleMock
+        sendAttributeSignal: expect.any( Function )
       } )
     );
     expect( getHandleMock ).toHaveBeenCalledWith( 'wf-1' );
   } );
 
+  it( 'handles next returning a non-Promise value', async () => {
+    const { ActivityExecutionInterceptor } = await import( './activity.js' );
+    const interceptor = new ActivityExecutionInterceptor( { activities: makeActivities(), workflows: makeWorkflows() } );
+    const next = vi.fn( () => ( { result: 'sync' } ) );
+
+    await expect( interceptor.execute( makeInput(), next ) ).resolves.toEqual( { result: 'sync' } );
+
+    expect( allSettledWithTimeoutMock ).toHaveBeenCalledWith( [], 30_000 );
+    expect( messageBusEmitMock ).toHaveBeenCalledWith( BusEventType.ACTIVITY_END, expect.any( Object ) );
+    expect( addEventEndMock ).toHaveBeenCalledWith( { id: 'act-1', details: { result: 'sync' }, executionContext: { workflowId: 'wf-1' } } );
+    expect( addEventErrorMock ).not.toHaveBeenCalled();
+  } );
+
+  it( 'handles signal flush timeout after successful execution', async () => {
+    const timeoutError = Object.assign( new Error( 'timeout' ), { isTimeout: true } );
+    allSettledWithTimeoutMock.mockRejectedValueOnce( timeoutError );
+    const { ActivityExecutionInterceptor } = await import( './activity.js' );
+    const interceptor = new ActivityExecutionInterceptor( { activities: makeActivities(), workflows: makeWorkflows() } );
+    const next = vi.fn().mockResolvedValue( { result: 'ok' } );
+
+    await expect( interceptor.execute( makeInput(), next ) ).resolves.toEqual( { result: 'ok' } );
+
+    expect( allSettledWithTimeoutMock ).toHaveBeenCalledWith( [], 30_000 );
+    expect( logWarnMock ).toHaveBeenCalledWith(
+      'Some usage/cost attributes were missed because not all activity signals were sent to the workflow',
+      { workflowId: 'wf-1', workflowName: 'myWorkflow', activityId: 'act-1', activityName: 'myWorkflow#myStep' }
+    );
+    expect( messageBusEmitMock ).toHaveBeenCalledWith( BusEventType.ACTIVITY_END, expect.any( Object ) );
+    expect( addEventEndMock ).toHaveBeenCalledOnce();
+    expect( addEventErrorMock ).not.toHaveBeenCalled();
+  } );
+
+  it( 'handles signal flush timeout after failed execution', async () => {
+    const timeoutError = Object.assign( new Error( 'timeout' ), { isTimeout: true } );
+    allSettledWithTimeoutMock.mockRejectedValueOnce( timeoutError );
+    const { ActivityExecutionInterceptor } = await import( './activity.js' );
+    const interceptor = new ActivityExecutionInterceptor( { activities: makeActivities(), workflows: makeWorkflows() } );
+    const error = new Error( 'step failed' );
+    const next = vi.fn().mockRejectedValue( error );
+
+    await expect( interceptor.execute( makeInput(), next ) ).rejects.toThrow( 'step failed' );
+
+    expect( allSettledWithTimeoutMock ).toHaveBeenCalledWith( [], 30_000 );
+    expect( logWarnMock ).toHaveBeenCalledWith(
+      'Some usage/cost attributes were missed because not all activity signals were sent to the workflow',
+      { workflowId: 'wf-1', workflowName: 'myWorkflow', activityId: 'act-1', activityName: 'myWorkflow#myStep' }
+    );
+    expect( messageBusEmitMock ).toHaveBeenCalledWith( BusEventType.ACTIVITY_ERROR, expect.objectContaining( { error } ) );
+    expect( addEventErrorMock ).toHaveBeenCalledOnce();
+    expect( addEventEndMock ).not.toHaveBeenCalled();
+  } );
+
+  it( 'exposes sendAttributeSignal in activity context', async () => {
+    const attribute = { setActivity: vi.fn() };
+    runWithContextMock.mockImplementationOnce( async ( fn, ctx ) => {
+      ctx.sendAttributeSignal( attribute );
+      return fn();
+    } );
+    const { ActivityExecutionInterceptor } = await import( './activity.js' );
+    const interceptor = new ActivityExecutionInterceptor( { activities: makeActivities(), workflows: makeWorkflows() } );
+    const next = vi.fn().mockResolvedValue( { result: 'ok' } );
+
+    await expect( interceptor.execute( makeInput(), next ) ).resolves.toEqual( { result: 'ok' } );
+
+    expect( attribute.setActivity ).toHaveBeenCalledWith( 'act-1', 'myWorkflow#myStep' );
+    expect( workflowHandleMock.signal ).toHaveBeenCalledWith( Signal.ADD_ATTRIBUTE, attribute );
+    expect( allSettledWithTimeoutMock ).toHaveBeenCalledWith( [ expect.any( Promise ) ], 30_000 );
+  } );
+
   it( 'records trace error event on failed execution', async () => {
     const { ActivityExecutionInterceptor } = await import( './activity.js' );
     const interceptor = new ActivityExecutionInterceptor( { activities: makeActivities(), workflows: makeWorkflows() } );