@output.ai/core 0.1.0 → 0.1.2

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

@@ -0,0 +1,67 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+vi.mock( '#utils', () => ( {
+  throws: e => {
+    throw e;
+  }
+} ) );
+
+const sendMock = vi.fn();
+const ctorState = { args: null };
+class S3ClientMock {
+  constructor( args ) {
+    ctorState.args = args;
+  } send = sendMock;
+}
+class PutObjectCommandMock {
+  constructor( input ) {
+    this.input = input;
+  }
+}
+
+vi.mock( '@aws-sdk/client-s3', () => ( {
+  S3Client: S3ClientMock,
+  PutObjectCommand: PutObjectCommandMock
+} ) );
+
+async function loadModule() {
+  vi.resetModules();
+  return import( './s3_client.js' );
+}
+
+describe( 'tracing/processors/s3/s3_client', () => {
+  beforeEach( () => {
+    vi.clearAllMocks();
+    delete process.env.AWS_REGION;
+    delete process.env.AWS_SECRET_ACCESS_KEY;
+    delete process.env.AWS_ACCESS_KEY_ID;
+    delete process.env.TRACE_REMOTE_S3_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.AWS_REGION = 'us-east-1';
+    process.env.AWS_SECRET_ACCESS_KEY = 'sek';
+    process.env.AWS_ACCESS_KEY_ID = 'id';
+    process.env.TRACE_REMOTE_S3_BUCKET = 'bucket';
+
+    const { upload } = await loadModule();
+
+    await upload( { key: 'wf/key.json', content: '{"a":1}' } );
+
+    expect( ctorState.args ).toEqual( { region: 'us-east-1', secretAccessKey: 'sek', accessKeyId: 'id' } );
+    expect( sendMock ).toHaveBeenCalledTimes( 1 );
+    const cmd = sendMock.mock.calls[0][0];
+    expect( cmd ).toBeInstanceOf( PutObjectCommandMock );
+    expect( cmd.input ).toEqual( { Bucket: 'bucket', Key: 'wf/key.json', Body: '{"a":1}' } );
+
+    // subsequent upload uses cached client
+    await upload( { key: 'wf/key2.json', content: '{}' } );
+    expect( sendMock ).toHaveBeenCalledTimes( 2 );
+  } );
+} );
+

package/src/tracing/{tracer_tree.js → tools/build_trace_tree.js}

@@ -1,6 +1,3 @@
-import { readFileSync, writeFileSync } from 'node:fs';
-import { EOL } from 'os';
-
 /**
  * @typedef {object} NodeEntry
  * @property {string} id
@@ -34,7 +31,7 @@ const createEntry = id => ( {
 } );
 
 /**
- * Build a tree of nodes from the raw trace file.
+ * Build a tree of nodes from a list of entries
  *
  * Each node will have: id, name, kind, children, input, output or error, startedAt, endedAt.
  *
@@ -48,14 +45,10 @@ const createEntry = id => ( {
  * Children are added according to the parentId of each entry.
  * The result tree has a single root: the only node without parentId, normally the workflow itself.
  *
- * @param {string} src - Full path to the raw log file (NDJSON)
+ * @param {object[]} entries - The list of entries
  * @returns {void}
  */
-export const buildLogTree = src => {
-  const content = readFileSync( src, 'utf-8' );
-  const lines = content.split( EOL ).filter( l => l.trim().length > 0 );
-  const entries = lines.map( l => JSON.parse( l ) );
-
+export default entries => {
   const nodes = new Map();
   const ensureNode = id => nodes.get( id ) ?? nodes.set( id, createEntry( id ) ).get( id );
 
@@ -79,5 +72,5 @@ export const buildLogTree = src => {
   }
 
   const root = nodes.get( entries.find( e => !e.parentId ).id );
-  writeFileSync( src.replace( /\.raw$/, '.json' ), JSON.stringify( root, undefined, 2 ), 'utf-8' );
+  return root;
 };

package/src/tracing/{tracer_tree.spec.js → tools/build_trace_tree.spec.js}

@@ -1,18 +1,8 @@
 import { describe, it, expect } from 'vitest';
-import { writeFileSync, readFileSync, rmSync } from 'node:fs';
-import { mkdtempSync } from 'node:fs';
-import { tmpdir } from 'node:os';
-import { join } from 'path';
-import { EOL } from 'os';
-import { buildLogTree } from './tracer_tree.js';
+import buildLogTree from './build_trace_tree.js';
 
-const createTempDir = () => mkdtempSync( join( tmpdir(), 'output-sdk-trace-tree-' ) );
-
-describe( 'tracer/tracer_tree', () => {
+describe( 'build_trace_tree', () => {
   it( 'builds a tree from workflow/step/IO entries with grouping and sorting', () => {
-    const tmp = createTempDir();
-    const rawPath = join( tmp, 'run-123.raw' );
-
     const entries = [
       // workflow start
       { kind: 'workflow', phase: 'start', name: 'wf', id: 'wf', parentId: undefined, details: { a: 1 }, timestamp: 1000 },
@@ -38,11 +28,7 @@ describe( 'tracer/tracer_tree', () => {
       { kind: 'workflow', phase: 'end', name: 'wf', id: 'wf', parentId: undefined, details: { ok: true }, timestamp: 3000 }
     ];
 
-    writeFileSync( rawPath, entries.map( e => JSON.stringify( e ) ).join( EOL ) + EOL, 'utf-8' );
-
-    buildLogTree( rawPath );
-
-    const jsonText = readFileSync( rawPath.replace( /.raw$/, '.json' ), 'utf-8' );
+    const result = buildLogTree( entries );
 
     const expected = {
       id: 'wf',
@@ -108,8 +94,6 @@ describe( 'tracer/tracer_tree', () => {
       ]
     };
 
-    expect( jsonText ).toBe( JSON.stringify( expected, undefined, 2 ) );
-
-    rmSync( tmp, { recursive: true, force: true } );
+    expect( result ).toMatchObject( expected );
   } );
 } );

package/src/tracing/{utils.js → tools/utils.js}

@@ -19,3 +19,10 @@ export const serializeError = error =>
     message: error.message,
     stack: error.stack
   };
+
+/**
+ * Returns true if string value is stringbool and true
+ * @param {string} v
+ * @returns
+ */
+export const isStringboolTrue = v => [ '1', 'true', 'on' ].includes( v );

package/src/tracing/trace_engine.js

@@ -0,0 +1,63 @@
+import { Storage } from '#async_storage';
+import { EventEmitter } from 'node:events';
+import { serializeError, isStringboolTrue } from './tools/utils.js';
+import * as localProcessor from './processors/local/index.js';
+import * as s3Processor from './processors/s3/index.js';
+
+const traceBus = new EventEmitter();
+const processors = [
+  {
+    isOn: isStringboolTrue( process.env.TRACE_LOCAL_ON ),
+    init: localProcessor.init,
+    exec: localProcessor.exec
+  },
+  {
+    isOn: isStringboolTrue( process.env.TRACE_REMOTE_ON ),
+    init: s3Processor.init,
+    exec: s3Processor.exec
+  }
+];
+
+/**
+ * Starts processors based on env vars and attach them to the main bus to listen trace events
+ */
+export const init = async () => {
+  for ( const p of processors.filter( p => p.isOn ) ) {
+    await p.init();
+    traceBus.addListener( 'entry', p.exec );
+  }
+};
+
+/**
+ * Serialize details of an event
+ */
+const serializeDetails = details => details instanceof Error ? serializeError( details ) : details;
+
+/**
+ * Creates a new trace event phase and sens to be written
+ *
+ * @param {string} phase - The phase
+ * @param {object} fields - All the trace fields
+ * @returns {void}
+ */
+export const addEventPhase = ( phase, { kind, name, id, parentId, details, executionContext } ) =>
+  traceBus.emit( 'entry', {
+    executionContext,
+    entry: { kind, phase, name, id, parentId, phase, timestamp: Date.now(), details: serializeDetails( details ) }
+  } );
+
+/**
+ * Adds an Event Phase, complementing the options with parentId and executionContext from the async storage.
+ *
+ * This function will have no effect if called from outside an Temporal Workflow/Activity environment,
+ * so it is safe to be used on unit tests or any dependencies that might be used elsewhere
+ *
+ * @param {object} options - The common trace configurations
+ */
+export function addEventPhaseWithContext( phase, options ) {
+  const storeContent = Storage.load();
+  if ( storeContent ) { // If there is no storageContext this was not called from an Temporal Environment
+    const { parentId, executionContext } = storeContent;
+    addEventPhase( phase, { ...options, parentId, executionContext } );
+  }
+};

package/src/tracing/trace_engine.spec.js

@@ -0,0 +1,91 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+const storageLoadMock = vi.fn();
+vi.mock( '#async_storage', () => ( {
+  Storage: { load: storageLoadMock }
+} ) );
+
+const localInitMock = vi.fn( async () => {} );
+const localExecMock = vi.fn();
+vi.mock( './processors/local/index.js', () => ( {
+  init: localInitMock,
+  exec: localExecMock
+} ) );
+
+const s3InitMock = vi.fn( async () => {} );
+const s3ExecMock = vi.fn();
+vi.mock( './processors/s3/index.js', () => ( {
+  init: s3InitMock,
+  exec: s3ExecMock
+} ) );
+
+async function loadTraceEngine() {
+  vi.resetModules();
+  return import( './trace_engine.js' );
+}
+
+describe( 'tracing/trace_engine', () => {
+  beforeEach( () => {
+    vi.clearAllMocks();
+    delete process.env.TRACE_LOCAL_ON;
+    delete process.env.TRACE_REMOTE_ON;
+    storageLoadMock.mockReset();
+  } );
+
+  it( 'init() starts only enabled processors and attaches listeners', async () => {
+    process.env.TRACE_LOCAL_ON = '1';
+    process.env.TRACE_REMOTE_ON = '0';
+    const { init, addEventPhase } = await loadTraceEngine();
+
+    await init();
+
+    expect( localInitMock ).toHaveBeenCalledTimes( 1 );
+    expect( s3InitMock ).not.toHaveBeenCalled();
+
+    addEventPhase( 'start', { kind: 'step', name: 'N', id: '1', parentId: 'p', details: { ok: true } } );
+    expect( localExecMock ).toHaveBeenCalledTimes( 1 );
+    const payload = localExecMock.mock.calls[0][0];
+    expect( payload.entry.name ).toBe( 'N' );
+    expect( payload.entry.kind ).toBe( 'step' );
+    expect( payload.entry.phase ).toBe( 'start' );
+    expect( payload.entry.details ).toEqual( { ok: true } );
+  } );
+
+  it( 'addEventPhase() emits an entry consumed by processors', async () => {
+    process.env.TRACE_LOCAL_ON = 'on';
+    const { init, addEventPhase } = await loadTraceEngine();
+    await init();
+
+    addEventPhase( 'end', { kind: 'workflow', name: 'W', id: '2', parentId: 'p2', details: 'done' } );
+    expect( localExecMock ).toHaveBeenCalledTimes( 1 );
+    const payload = localExecMock.mock.calls[0][0];
+    expect( payload.entry.name ).toBe( 'W' );
+    expect( payload.entry.phase ).toBe( 'end' );
+    expect( payload.entry.details ).toBe( 'done' );
+  } );
+
+  it( 'addEventPhaseWithContext() uses storage when available', async () => {
+    process.env.TRACE_LOCAL_ON = 'true';
+    storageLoadMock.mockReturnValue( { parentId: 'ctx-p', executionContext: { runId: 'r1' } } );
+    const { init, addEventPhaseWithContext } = await loadTraceEngine();
+    await init();
+
+    addEventPhaseWithContext( 'tick', { kind: 'step', name: 'S', id: '3', details: 1 } );
+    expect( localExecMock ).toHaveBeenCalledTimes( 1 );
+    const payload = localExecMock.mock.calls[0][0];
+    expect( payload.executionContext ).toEqual( { runId: 'r1' } );
+    expect( payload.entry.parentId ).toBe( 'ctx-p' );
+    expect( payload.entry.name ).toBe( 'S' );
+    expect( payload.entry.phase ).toBe( 'tick' );
+  } );
+
+  it( 'addEventPhaseWithContext() is a no-op when storage is absent', async () => {
+    process.env.TRACE_LOCAL_ON = '1';
+    storageLoadMock.mockReturnValue( undefined );
+    const { init, addEventPhaseWithContext } = await loadTraceEngine();
+    await init();
+
+    addEventPhaseWithContext( 'noop', { kind: 'step', name: 'X', id: '4', details: null } );
+    expect( localExecMock ).not.toHaveBeenCalled();
+  } );
+} );

package/src/utils.js

@@ -0,0 +1,37 @@
+import { METADATA_ACCESS_SYMBOL } from '#consts';
+
+/**
+ * Node safe clone implementation that doesn't use global structuredClone()
+ * @param {object} v
+ * @returns {object}
+ */
+export const clone = v => JSON.parse( JSON.stringify( v ) );
+
+/**
+ * Throw given error
+ * @param {Error} e
+ * @throws {e}
+ */
+export const throws = e => {
+  throw e;
+};
+
+/**
+ * Add metadata "values" property to a given object
+ * @param {object} target
+ * @param {object} values
+ * @returns
+ */
+export const setMetadata = ( target, values ) =>
+  Object.defineProperty( target, METADATA_ACCESS_SYMBOL, { value: values, writable: false, enumerable: false, configurable: false } );
+
+/**
+ * Merge two temporal activity options
+ * @param {import('@temporalio/workflow').ActivityOptions} base
+ * @param {import('@temporalio/workflow').ActivityOptions} ext
+ * @returns {import('@temporalio/workflow').ActivityOptions}
+ */
+export const mergeActivityOptions = ( base = {}, ext = {} ) =>
+  Object.entries( ext ).reduce( ( options, [ k, v ] ) =>
+    Object.assign( options, { [k]: typeof v === 'object' ? mergeActivityOptions( options[k], v ) : v } )
+  , clone( base ) );

package/src/utils.spec.js

@@ -0,0 +1,60 @@
+import { describe, it, expect } from 'vitest';
+import { clone, mergeActivityOptions } from './utils.js';
+
+describe( 'clone', () => {
+  it( 'produces a deep copy without shared references', () => {
+    const original = { a: 1, nested: { b: 2 } };
+    const copied = clone( original );
+
+    copied.nested.b = 3;
+
+    expect( original.nested.b ).toBe( 2 );
+    expect( copied.nested.b ).toBe( 3 );
+    expect( copied ).not.toBe( original );
+  } );
+} );
+
+describe( 'mergeActivityOptions', () => {
+  it( 'recursively merges nested objects', () => {
+    const base = {
+      taskQueue: 'q1',
+      retry: { maximumAttempts: 3, backoffCoefficient: 2 }
+    };
+    const ext = {
+      retry: { maximumAttempts: 5, initialInterval: '1s' }
+    };
+
+    const result = mergeActivityOptions( base, ext );
+
+    expect( result ).toEqual( {
+      taskQueue: 'q1',
+      retry: { maximumAttempts: 5, backoffCoefficient: 2, initialInterval: '1s' }
+    } );
+  } );
+
+  it( 'omitted properties in second do not overwrite first', () => {
+    const base = {
+      taskQueue: 'q2',
+      retry: { initialInterval: '2s', backoffCoefficient: 2 }
+    };
+    const ext = {
+      retry: { backoffCoefficient: 3 }
+    };
+
+    const result = mergeActivityOptions( base, ext );
+
+    expect( result.retry.initialInterval ).toBe( '2s' );
+    expect( result.retry.backoffCoefficient ).toBe( 3 );
+    expect( result.taskQueue ).toBe( 'q2' );
+  } );
+
+  it( 'handles omitted second argument by returning a clone', () => {
+    const base = { taskQueue: 'q3', retry: { maximumAttempts: 2 } };
+
+    const result = mergeActivityOptions( base );
+
+    expect( result ).toEqual( base );
+    expect( result ).not.toBe( base );
+  } );
+} );
+

package/src/worker/catalog_workflow/index.js

@@ -1,4 +1,5 @@
 import { z } from 'zod';
+import { dirname } from 'node:path';
 import { METADATA_ACCESS_SYMBOL } from '#consts';
 import { Catalog, CatalogActivity, CatalogWorkflow } from './catalog.js';
 
@@ -37,7 +38,7 @@ export const createCatalog = ( { workflows, activities } ) =>
       inputSchema: convertToJsonSchema( workflow.inputSchema ),
       outputSchema: convertToJsonSchema( workflow.outputSchema ),
       activities: Object.entries( activities )
-        .filter( ( [ k ] ) => k.startsWith( `${workflow.path}#` ) )
+        .filter( ( [ k ] ) => k.startsWith( `${dirname( workflow.path )}#` ) )
         .map( ( [ _, v ] ) => {
           const metadata = v[METADATA_ACCESS_SYMBOL];
           return new CatalogActivity( {

package/src/worker/catalog_workflow/index.spec.js

@@ -17,16 +17,14 @@ describe( 'createCatalog', () => {
     const workflows = [
       {
         name: 'flow1',
-        path: '/flows/flow1',
-        pathname: '/flows/flow1/workflow.js',
+        path: '/flows/flow1/workflow.js',
         description: 'desc-flow1',
         inputSchema: z.object( { in: z.literal( 'f1' ) } ),
         outputSchema: z.object( { out: z.literal( 'f1' ) } )
       },
       {
         name: 'flow2',
-        path: '/flows/flow2',
-        pathname: '/flows/flow2/workflow.js',
+        path: '/flows/flow2/workflow.js',
         description: 'desc-flow2',
         inputSchema: z.object( { in: z.literal( 'f2' ) } ),
         outputSchema: z.object( { out: z.literal( 'f2' ) } )
@@ -95,7 +93,7 @@ describe( 'createCatalog', () => {
     expect( mapped ).toEqual( [
       {
         name: 'flow1',
-        path: '/flows/flow1',
+        path: '/flows/flow1/workflow.js',
         description: 'desc-flow1',
         inputSchema: {
           $schema: 'https://json-schema.org/draft/2020-12/schema',
@@ -152,7 +150,7 @@ describe( 'createCatalog', () => {
       },
       {
         name: 'flow2',
-        path: '/flows/flow2',
+        path: '/flows/flow2/workflow.js',
         description: 'desc-flow2',
         inputSchema: {
           $schema: 'https://json-schema.org/draft/2020-12/schema',
@@ -192,9 +190,7 @@ describe( 'createCatalog', () => {
     ] );
 
     // Original inputs are not mutated
-    expect( workflows[0].path ).toBe( '/flows/flow1' );
-    expect( workflows[0].pathname ).toBe( '/flows/flow1/workflow.js' );
-    expect( workflows[1].path ).toBe( '/flows/flow2' );
-    expect( workflows[1].pathname ).toBe( '/flows/flow2/workflow.js' );
+    expect( workflows[0].path ).toBe( '/flows/flow1/workflow.js' );
+    expect( workflows[1].path ).toBe( '/flows/flow2/workflow.js' );
   } );
 } );

package/src/worker/configs.js

@@ -0,0 +1,24 @@
+import * as z from 'zod';
+
+class InvalidEnvVarsErrors extends Error { }
+
+const envVarSchema = z.object( {
+  CATALOG_ID: z.string().regex( /^[a-z0-9_.@-]+$/i ),
+  TEMPORAL_ADDRESS: z.string().default( 'localhost:7233' ),
+  TEMPORAL_API_KEY: z.string().optional(),
+  TEMPORAL_NAMESPACE: z.string().optional().default( 'default' )
+} );
+
+const { data: envVars, error } = envVarSchema.safeParse( process.env );
+if ( error ) {
+  throw new InvalidEnvVarsErrors( z.prettifyError( error ) );
+}
+
+export const address = envVars.TEMPORAL_ADDRESS;
+export const apiKey = envVars.TEMPORAL_API_KEY;
+export const executionTimeout = '1m';
+export const maxActivities = 100;
+export const maxWorkflows = 100;
+export const namespace = envVars.TEMPORAL_NAMESPACE;
+export const taskQueue = envVars.CATALOG_ID;
+export const catalogId = envVars.CATALOG_ID;

package/src/worker/index.js

@@ -3,16 +3,16 @@ import { Client } from '@temporalio/client';
 import { WorkflowIdConflictPolicy } from '@temporalio/common';
 import { dirname, join } from 'path';
 import { fileURLToPath } from 'node:url';
-import { worker as workerConfig } from '#configs';
+import { address, apiKey, maxActivities, maxWorkflows, namespace, taskQueue, catalogId } from './configs.js';
 import { loadActivities, loadWorkflows, createWorkflowsEntryPoint } from './loader.js';
 import { ActivityExecutionInterceptor } from './interceptors/activity.js';
 import { sinks } from './sinks.js';
 import { createCatalog } from './catalog_workflow/index.js';
+import { init as initTracing } from '#tracing';
+import { WORKFLOW_CATALOG } from '#consts';
 
 const __dirname = dirname( fileURLToPath( import.meta.url ) );
 
-const { address, apiKey, maxActivities, maxWorkflows, namespace, taskQueue, catalogId } = workerConfig;
-
 // Get caller directory from command line arguments
 const callerDir = process.argv[2];
 
@@ -26,6 +26,9 @@ const callerDir = process.argv[2];
   console.log( '[Core]', 'Creating worker entry point...' );
   const workflowsPath = createWorkflowsEntryPoint( workflows );
 
+  console.log( '[Core]', 'Initializing tracing...' );
+  await initTracing();
+
   console.log( '[Core]', 'Creating workflows catalog...' );
   const catalog = createCatalog( { workflows, activities } );
 
@@ -69,7 +72,7 @@ const callerDir = process.argv[2];
   } );
 
   console.log( '[Core]', 'Starting catalog workflow...' );
-  await new Client( { connection, namespace } ).workflow.start( 'catalog', {
+  await new Client( { connection, namespace } ).workflow.start( WORKFLOW_CATALOG, {
     taskQueue,
     workflowId: catalogId, // use the name of the task queue as the catalog name, ensuring uniqueness
     workflowIdConflictPolicy: WorkflowIdConflictPolicy.TERMINATE_EXISTING,

package/src/worker/interceptors/activity.js

@@ -20,26 +20,19 @@ export class ActivityExecutionInterceptor {
 
   async execute( input, next ) {
     const { workflowExecution: { workflowId }, activityId, activityType } = Context.current().info;
-    const { traceId, traceHelm } = headersToObject( input.headers );
-    const { type: kind, skipTrace } = this.activities?.[activityType]?.[METADATA_ACCESS_SYMBOL];
+    const { executionContext } = headersToObject( input.headers );
+    const { type: kind } = this.activities?.[activityType]?.[METADATA_ACCESS_SYMBOL];
 
-    const traceContext = { kind, id: activityId, parentId: workflowId, name: activityType, traceId, traceHelm };
-
-    if ( !skipTrace ) {
-      addEventStart( { details: input.args[0], ...traceContext } );
-    }
+    const traceArguments = { kind, id: activityId, parentId: workflowId, name: activityType, executionContext };
+    addEventStart( { details: input.args[0], ...traceArguments } );
 
     // creates a context for the nested tracing
     try {
-      const output = await Storage.runWithContext( async _ => next( input ), { parentId: activityId, traceId, traceHelm } );
-      if ( !skipTrace ) {
-        addEventEnd( { details: output, ...traceContext } );
-      }
+      const output = await Storage.runWithContext( async _ => next( input ), { parentId: activityId, executionContext } );
+      addEventEnd( { details: output, ...traceArguments } );
       return output;
     } catch ( error ) {
-      if ( !skipTrace ) {
-        addEventError( { details: error, ...traceContext } );
-      }
+      addEventError( { details: error, ...traceArguments } );
       throw error;
     }
   }

package/src/worker/interceptors/workflow.js

@@ -1,6 +1,9 @@
 // THIS RUNS IN THE TEMPORAL'S SANDBOX ENVIRONMENT
 import { workflowInfo, proxySinks, ApplicationFailure } from '@temporalio/workflow';
 import { memoToHeaders } from '../sandboxed_utils.js';
+import { mergeActivityOptions } from '#utils';
+// this is a dynamic generated file with activity configs overwrites
+import stepOptions from '../temp/__activity_options.js';
 
 /*
   This is not an AI comment!
@@ -13,7 +16,13 @@ import { memoToHeaders } from '../sandboxed_utils.js';
 */
 class HeadersInjectionInterceptor {
   async scheduleActivity( input, next ) {
-    Object.assign( input.headers, memoToHeaders( workflowInfo().memo ) );
+    const memo = workflowInfo().memo ?? {};
+    Object.assign( input.headers, memoToHeaders( memo ) );
+    // apply per-invocation options passed as second argument by rewritten calls
+    const options = stepOptions[input.activityType];
+    if ( options ) {
+      input.options = mergeActivityOptions( memo.activityOptions, options );
+    }
     return next( input );
   }
 };
@@ -25,8 +34,7 @@ class WorkflowExecutionInterceptor {
     sinks.trace.addWorkflowEventStart( input.args[0] );
     try {
       const output = await next( input );
-      // if the workflow is root, the result wraps "output" and "trace" onto an object, this will hide the trace
-      sinks.trace.addWorkflowEventEnd( !workflowInfo().memo.parentId ? output.output : output );
+      sinks.trace.addWorkflowEventEnd( output );
       return output;
     } catch ( error ) {
       sinks.trace.addWorkflowEventError( error );