Changed: Publish a LiveKit model response's links once

Author: delner

js/src/instrumentation/plugins/livekit-plugin.test.ts

@@ -13,6 +13,7 @@ import { z } from "zod/v3";
 import { configureNode } from "../../node/config";
 import {
   Attachment,
+  SpanImpl,
   _exportsForTestingOnly as logging,
   initLogger,
 } from "../../logger";
@@ -413,6 +414,36 @@ describe("tracing", () => {
     });
   });
 
+  test("a model response publishes the links from its context once, and again for a late target", async () => {
+    const root = start("agent_session");
+    const turn = start("agent_turn", root);
+    const tool = (callId: string) =>
+      start("function_tool", turn, {
+        "lk.function_tool.name": "lookup_order",
+        "gen_ai.tool.call.id": callId,
+      }).end();
+    for (const callId of ["call-1", "call-2", "call-3"]) tool(callId);
+    const log = vi.spyOn(SpanImpl.prototype, "log");
+    const publications = () =>
+      log.mock.calls.filter(([event]) => event.metadata?.links).length;
+    start("llm_node", turn, {
+      "lk.pii.chat_ctx": JSON.stringify({
+        items: ["call-1", "call-2", "call-3", "call-4"].map((call_id) => ({
+          type: "function_call_output",
+          call_id,
+        })),
+      }),
+    }).end();
+    expect(publications()).toBe(1);
+    tool("call-4");
+    expect(publications()).toBe(2);
+    turn.end();
+    const { one, named } = await rows();
+    expect(
+      one("llm_response").metadata.links.map((l: any) => l.context.spanId),
+    ).toEqual(named("lookup_order").map((r) => r.span_id));
+  });
+
   test("speech is generated by every synthesis in its reply, and only its reply", async () => {
     const root = start("agent_session");
     const turn = start("agent_turn", root);

js/src/instrumentation/plugins/livekit-plugin.ts

@@ -629,7 +629,10 @@ class LiveKitTracing {
     const state = this.sessions.get(row.session);
     if (!row.owner) {
       row.span.log(endEvent(row, attrs, events));
-      if (state) correlate(row, identities, state.relationships);
+      if (state)
+        state.relationships.batch(row.span, () =>
+          correlate(row, identities, state.relationships),
+        );
       if (state && row.name === "user_speaking")
         row.span.log({
           metadata: {
@@ -1287,6 +1290,19 @@ class Relationships {
   private targets = new Map<string, Span[]>();
   private waiting = new Map<string, Map<Span, Relationship>>();
   private links = new Map<Span, Map<string, Link>>();
+  /** A source whose new links publish together, and whether it has any. */
+  private batched?: { source: Span; changed: boolean };
+
+  /** Publishes the links `fn` adds from `source` once, as one complete list. */
+  batch(source: Span, fn: () => void): void {
+    const batched = (this.batched = { source, changed: false });
+    try {
+      fn();
+    } finally {
+      this.batched = undefined;
+      if (batched.changed) this.publish(source);
+    }
+  }
 
   register(id: string, target: Span): void {
     const targets = this.targets.get(id) ?? [];
@@ -1337,9 +1353,15 @@ class Relationships {
       attributes: { "braintrust.relationship.type": relationship },
     });
     this.links.set(source, links);
+    if (this.batched?.source === source) this.batched.changed = true;
+    else this.publish(source);
+  }
+
+  private publish(source: Span): void {
     // Native rows have no top-level links field, and OTel ingestion does not keep
     // links yet. Metadata carries them until links are first-class.
-    source.log({ metadata: { links: [...links.values()] } });
+    const links = [...(this.links.get(source)?.values() ?? [])];
+    source.log({ metadata: { links } });
   }
 }