@outputai/core 0.4.1-dev.92bc2fb.0 → 0.4.1-next.43c9293.0

package/src/tracing/trace_attribute.js

@@ -0,0 +1,80 @@
+import Decimal from 'decimal.js';
+
+/**
+ * All attributes inherit from this
+ */
+export class BaseAttribute {
+  activityId;
+  activityName;
+  date = Date.now();
+  type;
+
+  constructor( type ) {
+    this.type = type;
+  }
+
+  setActivity( id, name ) {
+    this.activityId = id;
+    this.activityName = name;
+  }
+}
+
+class HTTPRequestCount extends BaseAttribute {
+  static TYPE = 'http:request:count';
+  url;
+  requestId;
+
+  constructor( url, requestId ) {
+    super( HTTPRequestCount.TYPE );
+    this.url = url;
+    this.requestId = requestId;
+  }
+}
+
+class HTTPRequestCost extends BaseAttribute {
+  static TYPE = 'http:request:cost';
+  url;
+  requestId;
+  total = 0;
+
+  constructor( url, requestId, total ) {
+    super( HTTPRequestCost.TYPE );
+    this.url = url;
+    this.requestId = requestId;
+    this.total = total;
+  }
+}
+
+class LLMUsage extends BaseAttribute {
+  static TYPE = 'llm:usage';
+  modelId;
+  usage = [];
+  total = 0;
+  tokensUsed = 0;
+
+  constructor( modelId ) {
+    super( LLMUsage.TYPE );
+    this.modelId = modelId;
+  }
+
+  addUsage( { type, ppm, amount } ) {
+    const total = Decimal( amount ).div( 1_000_000 ).mul( ppm ).toNumber();
+    this.usage.push( {
+      type,
+      ppm,
+      amount,
+      total
+    } );
+    this.total = Decimal( this.total ).add( total ).toNumber();
+    this.tokensUsed = Decimal( this.tokensUsed ).add( amount ).toNumber();
+  }
+}
+
+/**
+ * Types of ADD_ATTR attributes
+ */
+export const Attribute = {
+  LLMUsage,
+  HTTPRequestCost,
+  HTTPRequestCount
+};

package/src/tracing/trace_engine.js

@@ -4,8 +4,10 @@ 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 } from '#consts';
+import { ComponentType, Signal } from '#consts';
 import { createChildLogger } from '#logger';
+import { EventAction } from './trace_consts.js';
+import { BaseAttribute } from './trace_attribute.js';
 
 const log = createChildLogger( 'Tracing' );
 
@@ -91,7 +93,15 @@ 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, executionContext } = storeContent;
+    const { parentId, parentName, executionContext, workflowHandle } = 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` );
+      }
+      attribute.setActivity( parentId, parentName );
+      workflowHandle.signal( Signal.ADD_ATTRIBUTE, attribute );
+    }
     addEventAction( action, { ...options, parentId, executionContext } );
   }
 };

package/src/worker/index.js

@@ -69,7 +69,7 @@ const callerDir = process.argv[2];
     workflowsPath,
     activities,
     sinks,
-    interceptors: initInterceptors( { activities, workflows } ),
+    interceptors: initInterceptors( { activities, workflows, connection } ),
     maxConcurrentWorkflowTaskExecutions,
     maxConcurrentActivityTaskExecutions,
     maxCachedWorkflows,

package/src/worker/index.spec.js

@@ -123,7 +123,7 @@ describe( 'worker/index', () => {
       maxConcurrentActivityTaskPolls: configValues.maxConcurrentActivityTaskPolls,
       maxConcurrentWorkflowTaskPolls: configValues.maxConcurrentWorkflowTaskPolls
     } ) );
-    expect( initInterceptorsMock ).toHaveBeenCalledWith( { activities: {}, workflows: [] } );
+    expect( initInterceptorsMock ).toHaveBeenCalledWith( { activities: {}, workflows: [], connection: mockConnection } );
     expect( registerShutdownMock ).toHaveBeenCalledWith( { worker: mockWorker, log: mockLog } );
     expect( startCatalogMock ).toHaveBeenCalledWith( {
       connection: mockConnection,

package/src/worker/interceptors/activity.js

@@ -5,6 +5,7 @@ import { headersToObject } from '../sandboxed_utils.js';
 import { BusEventType, METADATA_ACCESS_SYMBOL } from '#consts';
 import { activityHeartbeatEnabled, activityHeartbeatIntervalMs } from '../configs.js';
 import { messageBus } from '#bus';
+import { Client } from '@temporalio/client';
 
 /*
   This interceptor wraps every activity execution with cross-cutting concerns:
@@ -23,7 +24,7 @@ import { messageBus } from '#bus';
   - Headers injected by the workflow interceptor (executionContext)
 */
 export class ActivityExecutionInterceptor {
-  constructor( { activities, workflows } ) {
+  constructor( { activities, workflows, connection } ) {
     this.activities = activities;
     this.workflowsMap = workflows.reduce( ( map, w ) => {
       map.set( w.name, w );
@@ -32,14 +33,19 @@ export class ActivityExecutionInterceptor {
       }
       return map;
     }, new Map() );
+    this.connection = connection;
   };
 
   async execute( input, next ) {
     const startDate = Date.now();
+    const client = new Client( { connection: this.connection } );
+
     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 workflowHandle = client.workflow.getHandle( workflowId );
+
     messageBus.emit( BusEventType.ACTIVITY_START, { id, name, kind, workflowId, workflowName } );
     Tracing.addEventStart( { id, name, kind, parentId: workflowId, details: input.args[0], executionContext } );
 
@@ -56,7 +62,8 @@ export class ActivityExecutionInterceptor {
       intervals.heartbeat = activityHeartbeatEnabled && setInterval( () => Context.current().heartbeat(), activityHeartbeatIntervalMs );
 
       // Wraps the execution with accessible metadata for the activity
-      const output = await Storage.runWithContext( async _ => next( input ), { parentId: id, executionContext, workflowFilename } );
+      const ctx = { parentId: id, parentName: name, executionContext, workflowFilename, workflowHandle };
+      const output = await Storage.runWithContext( async _ => next( input ), ctx );
 
       messageBus.emit( BusEventType.ACTIVITY_END, { id, name, kind, workflowId, workflowName, duration: Date.now() - startDate } );
       Tracing.addEventEnd( { id, details: output, executionContext } );

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

@@ -2,6 +2,8 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { BusEventType } 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 heartbeatMock = vi.fn();
 const runWithContextMock = vi.hoisted( () => vi.fn().mockImplementation( async fn => fn() ) );
@@ -21,6 +23,14 @@ vi.mock( '@temporalio/activity', () => ( {
   }
 } ) );
 
+vi.mock( '@temporalio/client', () => ( {
+  Client: class Client {
+    workflow = {
+      getHandle: getHandleMock
+    };
+  }
+} ) );
+
 vi.mock( '#async_storage', () => ( {
   Storage: {
     runWithContext: runWithContextMock
@@ -108,12 +118,15 @@ describe( 'ActivityExecutionInterceptor', () => {
     expect( addEventErrorMock ).not.toHaveBeenCalled();
     expect( runWithContextMock ).toHaveBeenCalledWith(
       expect.any( Function ),
-      {
+      expect.objectContaining( {
         parentId: 'act-1',
+        parentName: 'myWorkflow#myStep',
         executionContext: { workflowId: 'wf-1' },
-        workflowFilename: '/workflows/myWorkflow.js'
-      }
+        workflowFilename: '/workflows/myWorkflow.js',
+        workflowHandle: workflowHandleMock
+      } )
     );
+    expect( getHandleMock ).toHaveBeenCalledWith( 'wf-1' );
   } );
 
   it( 'records trace error event on failed execution', async () => {

package/src/worker/interceptors.js

@@ -4,7 +4,7 @@ import { ActivityExecutionInterceptor } from './interceptors/activity.js';
 
 const __dirname = dirname( fileURLToPath( import.meta.url ) );
 
-export const initInterceptors = ( { activities, workflows } ) => ( {
+export const initInterceptors = ( { activities, workflows, connection } ) => ( {
   workflowModules: [ join( __dirname, './interceptors/workflow.js' ) ],
-  activityInbound: [ () => new ActivityExecutionInterceptor( { activities, workflows } ) ]
+  activityInbound: [ () => new ActivityExecutionInterceptor( { activities, workflows, connection } ) ]
 } );

package/src/tracing/tools/aggregate_trace_attributes.js

@@ -1,118 +0,0 @@
-/**
- * Aggregate `attributes.cost` and `attributes.token_usage` across an entire trace tree.
- *
- * Walks every node in the tree, sums `attributes.cost.total` grouped by the emitting
- * event name (inferred from node `kind` — see `eventNameForKind`), and sums
- * `attributes.token_usage` across LLM nodes. Falls back to `output.usage` on
- * legacy llm trace nodes that predate the `attributes.token_usage` write
- * (see overview §1.2).
- *
- * @typedef {object} TraceAttributes
- * @property {{ total: number, components: Array<{ name: string, value: number }> }} cost
- * @property {{ inputTokens: number, outputTokens: number, cachedInputTokens: number, totalTokens: number }} tokenUsage
- */
-
-const COST_EVENT_LLM = 'cost:llm:request';
-const COST_EVENT_HTTP = 'cost:http:request';
-const COST_EVENT_OTHER = 'other';
-
-/**
- * Map a trace node `kind` to the canonical cost event name that would emit it.
- * Unknown kinds bucket into `other` so future event sources still roll up cleanly.
- *
- * @param {string} kind
- * @returns {string}
- */
-const eventNameForKind = kind => {
-  if ( kind === 'llm' ) {
-    return COST_EVENT_LLM;
-  }
-  if ( kind === 'http' ) {
-    return COST_EVENT_HTTP;
-  }
-  return COST_EVENT_OTHER;
-};
-
-const isNumber = value => typeof value === 'number' && Number.isFinite( value );
-
-/**
- * Pull token usage off an llm node, preferring the new attribute over the legacy
- * `output.usage` fallback. Returns `null` when neither shape is present.
- */
-const readTokenUsage = node => {
-  const attrUsage = node.attributes?.token_usage;
-  if ( attrUsage && typeof attrUsage === 'object' ) {
-    return attrUsage;
-  }
-  const legacyUsage = node.output?.usage;
-  if ( legacyUsage && typeof legacyUsage === 'object' ) {
-    return legacyUsage;
-  }
-  return null;
-};
-
-/**
- * Recursively walk a trace tree depth-first, applying `visit` to each node.
- */
-const walk = ( node, visit ) => {
-  if ( !node ) {
-    return;
-  }
-  visit( node );
-  for ( const child of node.children ?? [] ) {
-    walk( child, visit );
-  }
-};
-
-/**
- * Build the aggregated `attributes` payload returned by `/trace-attributes`.
- * Component buckets always appear in a stable order so callers can index them
- * positionally if they want to.
- *
- * @param {object|null} root - The root NodeEntry returned by `buildTraceTree`.
- * @returns {TraceAttributes}
- */
-export default function aggregateTraceAttributes( root ) {
-  const costByEvent = new Map( [
-    [ COST_EVENT_LLM, 0 ],
-    [ COST_EVENT_HTTP, 0 ],
-    [ COST_EVENT_OTHER, 0 ]
-  ] );
-  const tokenUsage = { inputTokens: 0, outputTokens: 0, cachedInputTokens: 0, totalTokens: 0 };
-
-  walk( root, node => {
-    const cost = node.attributes?.cost;
-    if ( cost && isNumber( cost.total ) ) {
-      const eventName = eventNameForKind( node.kind );
-      costByEvent.set( eventName, ( costByEvent.get( eventName ) ?? 0 ) + cost.total );
-    }
-
-    if ( node.kind === 'llm' ) {
-      const usage = readTokenUsage( node );
-      if ( usage ) {
-        if ( isNumber( usage.inputTokens ) ) {
-          tokenUsage.inputTokens += usage.inputTokens;
-        }
-        if ( isNumber( usage.outputTokens ) ) {
-          tokenUsage.outputTokens += usage.outputTokens;
-        }
-        if ( isNumber( usage.cachedInputTokens ) ) {
-          tokenUsage.cachedInputTokens += usage.cachedInputTokens;
-        }
-        if ( isNumber( usage.totalTokens ) ) {
-          tokenUsage.totalTokens += usage.totalTokens;
-        }
-      }
-    }
-  } );
-
-  const components = Array.from( costByEvent, ( [ name, value ] ) => ( { name, value } ) );
-  const total = components.reduce( ( sum, { value } ) => sum + value, 0 );
-
-  return {
-    cost: { total, components },
-    tokenUsage
-  };
-}
-
-export { COST_EVENT_LLM, COST_EVENT_HTTP, COST_EVENT_OTHER };

package/src/tracing/tools/aggregate_trace_attributes.spec.js

@@ -1,231 +0,0 @@
-import { describe, it, expect } from 'vitest';
-import aggregateTraceAttributes, {
-  COST_EVENT_LLM,
-  COST_EVENT_HTTP,
-  COST_EVENT_OTHER
-} from './aggregate_trace_attributes.js';
-
-const node = ( { id, kind = 'step', attributes = {}, output, children = [] } ) => ( {
-  id,
-  kind,
-  name: id,
-  startedAt: 0,
-  endedAt: 0,
-  input: undefined,
-  output,
-  attributes,
-  children
-} );
-
-describe( 'aggregate_trace_attributes', () => {
-  it( 'returns zeros for a null root', () => {
-    const result = aggregateTraceAttributes( null );
-    expect( result.cost.total ).toBe( 0 );
-    expect( result.cost.components ).toEqual( [
-      { name: COST_EVENT_LLM, value: 0 },
-      { name: COST_EVENT_HTTP, value: 0 },
-      { name: COST_EVENT_OTHER, value: 0 }
-    ] );
-    expect( result.tokenUsage ).toEqual( {
-      inputTokens: 0, outputTokens: 0, cachedInputTokens: 0, totalTokens: 0
-    } );
-  } );
-
-  it( 'returns zeros for a tree with no cost or usage attributes', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [ node( { id: 's1' } ), node( { id: 's2' } ) ]
-    } );
-    const result = aggregateTraceAttributes( root );
-    expect( result.cost.total ).toBe( 0 );
-    expect( result.tokenUsage.totalTokens ).toBe( 0 );
-  } );
-
-  it( 'buckets cost by node kind into llm / http / other components', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [
-        node( { id: 'llm-1', kind: 'llm', attributes: { cost: { total: 0.20 } } } ),
-        node( { id: 'llm-2', kind: 'llm', attributes: { cost: { total: 0.10 } } } ),
-        node( { id: 'http-1', kind: 'http', attributes: { cost: { total: 0.50 } } } ),
-        // Unknown kind falls into the catch-all bucket
-        node( { id: 'step-1', kind: 'step', attributes: { cost: { total: 0.07 } } } )
-      ]
-    } );
-    const result = aggregateTraceAttributes( root );
-
-    const byName = Object.fromEntries( result.cost.components.map( c => [ c.name, c.value ] ) );
-    expect( byName[COST_EVENT_LLM] ).toBeCloseTo( 0.30, 10 );
-    expect( byName[COST_EVENT_HTTP] ).toBeCloseTo( 0.50, 10 );
-    expect( byName[COST_EVENT_OTHER] ).toBeCloseTo( 0.07, 10 );
-    expect( result.cost.total ).toBeCloseTo( 0.87, 10 );
-  } );
-
-  it( 'total equals the sum of all components', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [
-        node( { id: 'llm-1', kind: 'llm', attributes: { cost: { total: 0.1234 } } } ),
-        node( { id: 'http-1', kind: 'http', attributes: { cost: { total: 0.0011 } } } )
-      ]
-    } );
-    const { cost } = aggregateTraceAttributes( root );
-    const sum = cost.components.reduce( ( s, c ) => s + c.value, 0 );
-    expect( cost.total ).toBeCloseTo( sum, 10 );
-  } );
-
-  it( 'sums token_usage across llm nodes from the attribute path', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [
-        node( {
-          id: 'llm-1', kind: 'llm', attributes: {
-            token_usage: { inputTokens: 100, outputTokens: 20, cachedInputTokens: 5, totalTokens: 125 }
-          }
-        } ),
-        node( {
-          id: 'llm-2', kind: 'llm', attributes: {
-            token_usage: { inputTokens: 50, outputTokens: 10, cachedInputTokens: 1, totalTokens: 61 }
-          }
-        } )
-      ]
-    } );
-    const { tokenUsage } = aggregateTraceAttributes( root );
-    expect( tokenUsage ).toEqual( {
-      inputTokens: 150,
-      outputTokens: 30,
-      cachedInputTokens: 6,
-      totalTokens: 186
-    } );
-  } );
-
-  it( 'falls back to output.usage on legacy llm nodes that lack attributes.token_usage', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [
-        // Legacy shape — usage lives on output.usage, no attributes.token_usage
-        node( {
-          id: 'llm-legacy',
-          kind: 'llm',
-          output: { result: '...', usage: { inputTokens: 200, outputTokens: 40, totalTokens: 240 } }
-        } )
-      ]
-    } );
-    const { tokenUsage } = aggregateTraceAttributes( root );
-    expect( tokenUsage.inputTokens ).toBe( 200 );
-    expect( tokenUsage.outputTokens ).toBe( 40 );
-    expect( tokenUsage.totalTokens ).toBe( 240 );
-    expect( tokenUsage.cachedInputTokens ).toBe( 0 );
-  } );
-
-  it( 'prefers attributes.token_usage over output.usage when both are present', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [
-        node( {
-          id: 'llm-1',
-          kind: 'llm',
-          attributes: { token_usage: { inputTokens: 10, outputTokens: 2, totalTokens: 12 } },
-          output: { usage: { inputTokens: 999, outputTokens: 999, totalTokens: 999 } }
-        } )
-      ]
-    } );
-    const { tokenUsage } = aggregateTraceAttributes( root );
-    expect( tokenUsage.inputTokens ).toBe( 10 );
-    expect( tokenUsage.totalTokens ).toBe( 12 );
-  } );
-
-  it( 'ignores token_usage shapes on non-llm nodes', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      // attributes.token_usage on a non-llm node is intentionally ignored —
-      // only llm nodes contribute to the token-usage rollup today.
-      children: [
-        node( {
-          id: 'step-1', kind: 'step', attributes: {
-            token_usage: { inputTokens: 999, outputTokens: 999, totalTokens: 999 }
-          }
-        } )
-      ]
-    } );
-    const { tokenUsage } = aggregateTraceAttributes( root );
-    expect( tokenUsage.totalTokens ).toBe( 0 );
-  } );
-
-  it( 'aggregates a mixed tree with cost on http nodes and usage on llm nodes', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [
-        node( {
-          id: 'llm-1',
-          kind: 'llm',
-          attributes: {
-            cost: { total: 0.0038 },
-            token_usage: { inputTokens: 2264, outputTokens: 411, cachedInputTokens: 100, totalTokens: 2775 }
-          }
-        } ),
-        node( {
-          id: 'http-1',
-          kind: 'http',
-          attributes: { cost: { total: 0.50 } }
-        } )
-      ]
-    } );
-    const result = aggregateTraceAttributes( root );
-
-    const byName = Object.fromEntries( result.cost.components.map( c => [ c.name, c.value ] ) );
-    expect( byName[COST_EVENT_LLM] ).toBeCloseTo( 0.0038, 10 );
-    expect( byName[COST_EVENT_HTTP] ).toBeCloseTo( 0.50, 10 );
-    expect( byName[COST_EVENT_OTHER] ).toBe( 0 );
-    expect( result.cost.total ).toBeCloseTo( 0.5038, 10 );
-
-    expect( result.tokenUsage ).toEqual( {
-      inputTokens: 2264,
-      outputTokens: 411,
-      cachedInputTokens: 100,
-      totalTokens: 2775
-    } );
-  } );
-
-  it( 'recurses through nested children', () => {
-    const root = node( {
-      id: 'wf',
-      kind: 'workflow',
-      children: [
-        node( {
-          id: 's1',
-          kind: 'step',
-          children: [
-            node( {
-              id: 'llm-1', kind: 'llm', attributes: {
-                cost: { total: 0.01 },
-                token_usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 }
-              }
-            } )
-          ]
-        } )
-      ]
-    } );
-    const result = aggregateTraceAttributes( root );
-    expect( result.cost.total ).toBeCloseTo( 0.01, 10 );
-    expect( result.tokenUsage.totalTokens ).toBe( 15 );
-  } );
-
-  it( 'keeps the canonical component ordering: llm, http, other', () => {
-    const root = node( { id: 'wf', kind: 'workflow' } );
-    const { cost } = aggregateTraceAttributes( root );
-    expect( cost.components.map( c => c.name ) ).toEqual( [
-      COST_EVENT_LLM,
-      COST_EVENT_HTTP,
-      COST_EVENT_OTHER
-    ] );
-  } );
-} );

package/src/tracing/tools/index.js

@@ -1,7 +0,0 @@
-export { default as buildTraceTree } from './build_trace_tree.js';
-export {
-  default as aggregateTraceAttributes,
-  COST_EVENT_LLM,
-  COST_EVENT_HTTP,
-  COST_EVENT_OTHER
-} from './aggregate_trace_attributes.js';