From 12ad1ffb0577d96accf738804cc307b514a58989 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Fri, 2 Oct 2026 16:31:38 -0400 Subject: [PATCH 1/2] fix(client): give each graph() run its own run id graph() caches the built graph per context so repeated runs skip re-evaluating every node's flag. The build also created the graph's track data, run id included, so every run for the same context reported the same runId on its $ld:ai:graph:* events. Each run now copies the cached track data with a fresh runId and passes it to route and stream_route through their opts, so the node, handoff and graph-level events of one run share one id and the next run gets another. The flag evaluations stay cached. --- .../src/launchdarkly_ai_server/graph.py | 39 +++++++++++++------ packages/client/tests/test_graph_stream.py | 30 ++++++++++++++ 2 files changed, 57 insertions(+), 12 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 02e81f55..54935064 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -34,6 +34,11 @@ MAX_TRAVERSAL_DEPTH = 100 MAX_GRAPH_CACHE_SIZE = 512 +# graph() caches a built graph per context, track data included. Each run passes its own copy +# with a fresh runId through route/stream_route opts under this key, so the handoff events +# emitted inside the cached build carry the run's id instead of the build's. +_RUN_TRACK_DATA = "_run_track_data" + def _sanitize_name(key: str) -> str: # Match TS sanitizeName: hyphens become underscores (tool names must be [a-zA-Z0-9_]). @@ -259,7 +264,7 @@ async def run_node( "$ld:ai:graph:handoff_success", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": from_node.key, "targetKey": node.key, }, @@ -277,7 +282,7 @@ async def run_node( "$ld:ai:graph:handoff_failure", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": from_node.key, "targetKey": node.key, }, @@ -435,7 +440,7 @@ async def route( "$ld:ai:graph:handoff_success", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": node.key, "targetKey": next_node.key, }, @@ -455,7 +460,7 @@ async def route( "$ld:ai:graph:handoff_failure", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": node.key, "targetKey": chosen_key, }, @@ -551,7 +556,7 @@ async def stream_node( "$ld:ai:graph:handoff_success", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": from_node.key, "targetKey": node.key, }, @@ -580,7 +585,7 @@ async def stream_node( "$ld:ai:graph:handoff_failure", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": from_node.key, "targetKey": node.key, }, @@ -694,7 +699,7 @@ async def stream_route( "$ld:ai:graph:handoff_success", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": node.key, "targetKey": next_node.key, }, @@ -725,7 +730,7 @@ async def stream_route( "$ld:ai:graph:handoff_failure", ld_ctx, { - **graph_track_data, + **opts.get(_RUN_TRACK_DATA, graph_track_data), "sourceKey": node.key, "targetKey": chosen_key, }, @@ -907,11 +912,13 @@ async def invoke( # Evict an arbitrary entry to keep the cache bounded. self._cache.pop(next(iter(self._cache))) self._cache[cache_key] = built - graph_def, graph_track_data, _ = built + graph_def, built_track_data, _ = built if not graph_def.enabled: raise ValueError(f'Agent graph "{self._key}" is disabled') + graph_track_data: TrackData = {**built_track_data, "runId": str(uuid.uuid4())} + tracer = trace.get_tracer("@launchdarkly/ai-server") with tracer.start_as_current_span("launchdarkly.graph") as span: span.set_attribute("launchdarkly.graph.key", self._key) @@ -944,7 +951,10 @@ async def invoke( 1, ) entered += 1 - opts: dict[str, Any] = {"variables": variables} + opts: dict[str, Any] = { + "variables": variables, + _RUN_TRACK_DATA: graph_track_data, + } if previous_node: opts["from"] = previous_node # History seeds the entry point only. After the root hop, nodes @@ -1123,11 +1133,13 @@ async def _stream_events( if len(self._cache) >= MAX_GRAPH_CACHE_SIZE: self._cache.pop(next(iter(self._cache))) self._cache[cache_key] = built - graph_def, graph_track_data, stream_route = built + graph_def, built_track_data, stream_route = built if not graph_def.enabled: raise ValueError(f'Agent graph "{self._key}" is disabled') + graph_track_data: TrackData = {**built_track_data, "runId": str(uuid.uuid4())} + tracer = trace.get_tracer("@launchdarkly/ai-server") span = tracer.start_span("launchdarkly.graph", context=caller_context) span.set_attribute("launchdarkly.graph.key", self._key) @@ -1161,7 +1173,10 @@ async def _stream_events( 1, ) entered += 1 - route_opts: dict[str, Any] = {"variables": variables} + route_opts: dict[str, Any] = { + "variables": variables, + _RUN_TRACK_DATA: graph_track_data, + } if previous_node: route_opts["from"] = previous_node # History seeds the entry point only. After the root hop, nodes diff --git a/packages/client/tests/test_graph_stream.py b/packages/client/tests/test_graph_stream.py index 2b1cd47c..9b011aa6 100644 --- a/packages/client/tests/test_graph_stream.py +++ b/packages/client/tests/test_graph_stream.py @@ -755,6 +755,36 @@ async def invoke_fn( # --------------------------------------------------------------------------- +class TestGraphRunId: + @pytest.mark.parametrize("mode", ["invoke", "stream"]) + async def test_each_run_gets_its_own_run_id( + self, mock_ld_client: MagicMock, mode: str + ) -> None: + """One graph() instance, one context, two runs: one id per run, shared within it.""" + g = graph("graph-key", handlers=[_make_streaming_handler(["ok"])]) + + run_ids: list[set[str]] = [] + for _ in range(2): + mock_ld_client.track.reset_mock() + if mode == "invoke": + await g.invoke("hi", CONTEXT) + else: + await _collect(g.stream("hi", CONTEXT)) + graph_events = [ + c[0] + for c in mock_ld_client.track.call_args_list + if c[0][0].startswith("$ld:ai:graph:") + ] + names = {e[0] for e in graph_events} + assert "$ld:ai:graph:handoff_success" in names + assert "$ld:ai:graph:invocation_success" in names + run_ids.append({e[2]["runId"] for e in graph_events}) + + assert len(run_ids[0]) == 1 + assert len(run_ids[1]) == 1 + assert run_ids[0] != run_ids[1] + + class TestGraphStreamOtel: async def test_stamps_conversation_id_when_bound_at_call_time( self, mock_ld_client: MagicMock From 496540436858c163e4edd49c1f3a552865be1c2e Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 13 Oct 2026 10:14:29 -0700 Subject: [PATCH 2/2] test(graph): check the span's run id in the two-runs test Each run's launchdarkly.graph span must carry the runId of that run's graph-level events, for both invoke and stream, and the two runs' spans must carry different ids. --- packages/client/tests/test_graph_stream.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/packages/client/tests/test_graph_stream.py b/packages/client/tests/test_graph_stream.py index 47e57fe4..7efa47e6 100644 --- a/packages/client/tests/test_graph_stream.py +++ b/packages/client/tests/test_graph_stream.py @@ -760,12 +760,15 @@ class TestGraphRunId: async def test_each_run_gets_its_own_run_id( self, mock_ld_client: MagicMock, mode: str ) -> None: - """One graph() instance, one context, two runs: one id per run, shared within it.""" + """One graph() instance, one context, two runs: one id per run, shared by the + run's graph span and its graph-level events.""" g = graph("graph-key", handlers=[_make_streaming_handler(["ok"])]) run_ids: list[set[str]] = [] + span_run_ids: list[Any] = [] for _ in range(2): mock_ld_client.track.reset_mock() + _exporter.clear() if mode == "invoke": await g.invoke("hi", CONTEXT) else: @@ -779,10 +782,18 @@ async def test_each_run_gets_its_own_run_id( assert "$ld:ai:graph:handoff_success" in names assert "$ld:ai:graph:invocation_success" in names run_ids.append({e[2]["runId"] for e in graph_events}) + graph_spans = [s for s in _finished() if s.name == "launchdarkly.graph"] + assert len(graph_spans) == 1 + span_run_ids.append( + (graph_spans[0].attributes or {}).get("launchdarkly.run.id") + ) assert len(run_ids[0]) == 1 assert len(run_ids[1]) == 1 assert run_ids[0] != run_ids[1] + assert {span_run_ids[0]} == run_ids[0] + assert {span_run_ids[1]} == run_ids[1] + assert span_run_ids[0] != span_run_ids[1] class TestGraphStreamOtel: