fix: correct end_nanos accounting for async model calls

pdl-live-react/src/view/timeline/model.ts

@@ -75,12 +75,25 @@ function computeModelIter(
         block,
       }
 
-  return [
-    ...(ignoreRoot ? [] : [root]),
-    ...childrenOf(block)
-      .filter(nonNullable)
-      .flatMap((child) => computeModelIter(child, root)),
-  ].filter(nonNullable)
+  const childrenModel = childrenOf(block)
+    .filter(nonNullable)
+    .flatMap((child) => computeModelIter(child, root))
+
+  // Correct for anomalies in the trace where a child may have an
+  // earlier end timestamp than its children. See
+  // https://github.com/IBM/prompt-declaration-language/pull/683
+  if (root) {
+    const maxEnd = childrenModel.reduce(
+      (maxEnd, child) => Math.max(maxEnd, child.block.pdl__timing.end_nanos),
+      0,
+    )
+    root.block.pdl__timing.end_nanos = Math.max(
+      root.block.pdl__timing.end_nanos,
+      maxEnd,
+    )
+  }
+
+  return [...(ignoreRoot ? [] : [root]), ...childrenModel].filter(nonNullable)
 }
 
 export function childrenOf(block: NonScalarPdlBlock) {

src/pdl/pdl_dumper.py

@@ -97,8 +97,6 @@ def block_to_dict(  # noqa: C901
             context = block.context
         if len(context) > 0:
             d["context"] = context
-    if block.pdl__timing is not None:
-        d["pdl__timing"] = timing_to_dict(block.pdl__timing)
     if block.description is not None:
         d["description"] = block.description
     if block.role is not None:
@@ -239,6 +237,12 @@ def block_to_dict(  # noqa: C901
     #     d["location"] = location_to_dict(block.location)
     if block.fallback is not None:
         d["fallback"] = block_to_dict(block.fallback, json_compatible)
+    # Warning: remember to update timing here at the end! this ensures
+    # that any logic that updates timestamps when futures
+    # finish... has a chance to do its work before we record the
+    # timestamps to the trace
+    if block.pdl__timing is not None:
+        d["pdl__timing"] = timing_to_dict(block.pdl__timing)
     return d
 
 

src/pdl/pdl_llms.py

@@ -100,6 +100,16 @@ def generate_text(
         pdl_future: PdlLazy[tuple[dict[str, Any], Any]] = PdlConst(future)
         message = lazy_apply((lambda x: x[0]), pdl_future)
         response = lazy_apply((lambda x: x[1]), pdl_future)
+
+        # update the end timestamp when the future is done
+        def update_end_nanos(future):
+            import time
+
+            if block.pdl__timing is not None:
+                block.pdl__timing.end_nanos = time.time_ns()
+
+        future.add_done_callback(update_end_nanos)
+
         return message, response
 
     @staticmethod