diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 52790ce9..f41e15e2 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -40,6 +40,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_]). @@ -268,7 +273,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, }, @@ -286,7 +291,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, }, @@ -444,7 +449,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, }, @@ -464,7 +469,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, }, @@ -560,7 +565,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, }, @@ -589,7 +594,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, }, @@ -703,7 +708,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, }, @@ -734,7 +739,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, }, @@ -932,11 +937,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: set_ld_span_attributes( @@ -971,7 +978,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 @@ -1164,11 +1174,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) set_ld_span_attributes(span, {"__ld": graph_track_data, "ldContext": context}) @@ -1202,7 +1214,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 764dbef3..7efa47e6 100644 --- a/packages/client/tests/test_graph_stream.py +++ b/packages/client/tests/test_graph_stream.py @@ -755,6 +755,47 @@ 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 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: + 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}) + 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: async def test_stamps_conversation_id_when_bound_at_call_time( self, mock_ld_client: MagicMock