11import { AsyncLocalStorage , executionAsyncId } from 'node:async_hooks' ;
22import { AuthoredFlowExecutionError } from './authored-flow-error.js' ;
33import { AuthoredPromiseGraph } from './authored-promise-graph.js' ;
4+ import { AuthoredStepGraph } from './authored-step-graph.js' ;
5+ import type { AuthoredStepEdges } from './authored-step-index.js' ;
46
57type OperationToken = object ;
68type ResolverProbe = Set < number > ;
@@ -66,7 +68,7 @@ const observedCombinators: Record<CombinatorName, Combinator> = Object.fromEntri
6668 const members = collectMembers ( values ) ;
6769 if ( members === undefined ) return native . call ( this , values ) ;
6870 const aggregate = native . call ( this , members ) ;
69- activeLifecycle . getStore ( ) ?. registerCombinator ( members , aggregate ) ;
71+ activeLifecycle . getStore ( ) ?. registerCombinator ( members , aggregate , name ) ;
7072 return aggregate ;
7173 } ;
7274 Object . defineProperty ( observed , 'name' , { value : name , configurable : true } ) ;
@@ -117,6 +119,8 @@ export class AuthoredFlowLifecycle {
117119 */
118120 applyPredicateGate : ( < T > ( operation : { id : string ; predicateGate : unknown } , value : T ) => Promise < T > ) | undefined = undefined ;
119121 private readonly graph : AuthoredPromiseGraph ;
122+ private readonly stepGraph : AuthoredStepGraph ;
123+ private readonly operationPromises = new Map < OperationToken , number > ( ) ;
120124 private readonly activeResolverProbes : ResolverProbe [ ] = [ ] ;
121125 private readonly invocations = new Map < OperationToken , AuthoredOperationInvocation [ ] > ( ) ;
122126 private readonly promiseAllAggregates = new Map < OperationToken , Set < number > > ( ) ;
@@ -135,6 +139,7 @@ export class AuthoredFlowLifecycle {
135139 for ( const probe of this . activeResolverProbes ) probe . add ( asyncId ) ;
136140 } ,
137141 ) ;
142+ this . stepGraph = new AuthoredStepGraph ( ( asyncId ) => this . graph . causesOf ( asyncId ) ) ;
138143 installPromiseAllObserver ( ) ;
139144 try {
140145 this . graph . enable ( ) ;
@@ -152,16 +157,25 @@ export class AuthoredFlowLifecycle {
152157 stepOwners . set ( step , { lifecycle : this , operation } ) ;
153158 }
154159
155- registerCombinator ( values : readonly unknown [ ] , aggregate : Promise < unknown > ) : void {
160+ registerCombinator (
161+ values : readonly unknown [ ] ,
162+ aggregate : Promise < unknown > ,
163+ combinator : CombinatorName ,
164+ ) : void {
156165 const aggregateId = this . graph . idOf ( aggregate ) ;
157166 if ( aggregateId === undefined ) return ;
158167 this . graph . registerRoot ( aggregateId ) ;
159168 const memberIds = new Set < number > ( ) ;
169+ const awaitedMembers : number [ ] = [ ] ;
160170 for ( const value of values ) {
161171 if ( ( typeof value !== 'object' && typeof value !== 'function' ) || value === null ) continue ;
162172 const memberId = this . graph . idOf ( value ) ;
163173 if ( memberId !== undefined ) memberIds . add ( memberId ) ;
164174 const owner = stepOwners . get ( value ) ;
175+ // A Step is a thenable, not a promise: its member identity for the step
176+ // graph is the operation's own promise.
177+ const awaited = owner ?. lifecycle === this ? this . operationPromises . get ( owner . operation ) : memberId ;
178+ if ( awaited !== undefined ) awaitedMembers . push ( awaited ) ;
165179 if ( owner ?. lifecycle !== this ) continue ;
166180 let operationAggregates = this . promiseAllAggregates . get ( owner . operation ) ;
167181 if ( operationAggregates === undefined ) {
@@ -172,6 +186,33 @@ export class AuthoredFlowLifecycle {
172186 }
173187 this . promiseAllGroups . push ( { aggregate : aggregateId , members : memberIds } ) ;
174188 this . registrations ++ ;
189+ this . stepGraph . registerAggregate ( aggregateId , aggregate , combinator , awaitedMembers ) ;
190+ }
191+
192+ /**
193+ * Record an operation's own promise and, for a step the run's DAG draws, its
194+ * place in that DAG. Called synchronously from the operation's constructor:
195+ * the causal predecessors are read from the context that invoked
196+ * `f.agent`/`f.run`/`f.llm`, which is observable only at that moment.
197+ *
198+ * Every operation's promise is recorded so `Promise.all([step, …])` can
199+ * expand to it; only a `node` becomes a frontier the walk stops at. An
200+ * operation that is not drawn (a helper effect) is walked through instead,
201+ * so an edge never names a step the index has no record for.
202+ */
203+ registerOperation (
204+ operation : OperationToken ,
205+ promise : Promise < unknown > ,
206+ node ?: { readonly step : string ; readonly label ?: string } ,
207+ ) : AuthoredStepEdges | undefined {
208+ const promiseId = this . graph . idOf ( promise ) ;
209+ if ( promiseId !== undefined ) this . operationPromises . set ( operation , promiseId ) ;
210+ if ( node === undefined ) return undefined ;
211+ return this . stepGraph . registerStep ( node . step , promiseId , executionAsyncId ( ) , node . label ) ;
212+ }
213+
214+ stepEdges ( step : string ) : AuthoredStepEdges | undefined {
215+ return this . stepGraph . edgesOf ( step ) ;
175216 }
176217
177218 registerInvocation (
@@ -304,6 +345,8 @@ export class AuthoredFlowLifecycle {
304345 this . closed = true ;
305346 this . graph . disable ( ) ;
306347 this . graph . clear ( ) ;
348+ this . stepGraph . release ( ) ;
349+ this . operationPromises . clear ( ) ;
307350 this . invocations . clear ( ) ;
308351 this . promiseAllAggregates . clear ( ) ;
309352 this . promiseAllGroups . length = 0 ;
0 commit comments