From 12ad1ffb0577d96accf738804cc307b514a58989 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Fri, 2 Oct 2026 16:31:38 -0400 Subject: [PATCH] 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