From 4dd266368c7a9fc4a4e45999db808e448a518e56 Mon Sep 17 00:00:00 2001 From: Jeff Dupont Date: Tue, 20 Oct 2026 10:27:23 -0700 Subject: [PATCH] fix(AIC-3570): re-resolve graph flags on every invoke and stream graph() cached the built graph per context for the life of the instance, so flag changes and the graph kill switch never reached a reused instance. Co-Authored-By: Claude Opus 5.5 --- .../src/launchdarkly_ai_server/graph.py | 36 +------ packages/client/tests/test_graph.py | 101 ++++++++++++------ 2 files changed, 73 insertions(+), 64 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 46be8969..8b67f6d8 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -2,7 +2,6 @@ import asyncio import inspect as _inspect -import json import logging import re as _re import time @@ -39,11 +38,10 @@ logger = logging.getLogger(__name__) 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 +# The built graph's track data is minted before the run's runId. 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. +# emitted inside the build carry the run's id instead of the build's. _RUN_TRACK_DATA = "_run_track_data" @@ -878,9 +876,6 @@ def __init__( ) -> None: self._key = key self._options = options - self._cache: dict[ - str, tuple[GraphDefinition, TrackData, Callable[..., Any]] - ] = {} async def invoke( self, @@ -927,19 +922,7 @@ async def _invoke( "use resolve_graph() with a framework-native runner." ) - try: - cache_key: str | None = json.dumps(context, sort_keys=True) - except (TypeError, ValueError): - cache_key = None - if cache_key is not None and cache_key in self._cache: - built = self._cache[cache_key] - else: - built = await _build_graph(self._key, context, resolved_options) - if cache_key is not None: - if len(self._cache) >= MAX_GRAPH_CACHE_SIZE: - # Evict an arbitrary entry to keep the cache bounded. - self._cache.pop(next(iter(self._cache))) - self._cache[cache_key] = built + built = await _build_graph(self._key, context, resolved_options) graph_def, built_track_data, _ = built if not graph_def.enabled: @@ -1165,18 +1148,7 @@ async def _stream_events( "use resolve_graph() with a framework-native runner." ) - try: - cache_key: str | None = json.dumps(context, sort_keys=True) - except (TypeError, ValueError): - cache_key = None - if cache_key is not None and cache_key in self._cache: - built = self._cache[cache_key] - else: - built = await _build_graph(self._key, context, resolved_options) - if cache_key is not None: - if len(self._cache) >= MAX_GRAPH_CACHE_SIZE: - self._cache.pop(next(iter(self._cache))) - self._cache[cache_key] = built + built = await _build_graph(self._key, context, resolved_options) graph_def, built_track_data, stream_route = built if not graph_def.enabled: diff --git a/packages/client/tests/test_graph.py b/packages/client/tests/test_graph.py index c3379682..e1653ee2 100644 --- a/packages/client/tests/test_graph.py +++ b/packages/client/tests/test_graph.py @@ -324,14 +324,80 @@ async def test_graph_node_tracked_per_visited_node( assert node_events[0][2]["graphKey"] == "graph-key" assert "path" not in node_events[0][2] - async def test_node_variations_resolved_only_once( + async def test_node_variations_resolved_on_every_call( self, mock_ld_client: MagicMock ) -> None: + """Each invoke() re-resolves the graph and node flags. See TESTING.md §3.11.""" g = graph("graph-key", handlers=[_make_handler()]) await g.invoke("hi", CONTEXT) + after_first = mock_ld_client.variation.call_count await g.invoke("hi again", CONTEXT) - # Second call should use the cache — variation calls should not double - assert mock_ld_client.variation.call_count <= 4 # graph + 2 nodes × 1 cache hit + assert mock_ld_client.variation.call_count == after_first * 2 + + async def test_flag_change_reaches_reused_graph( + self, mock_ld_client: MagicMock + ) -> None: + """A node flag change between calls is seen on the next call. See TESTING.md §3.11.""" + state = {"instructions": "v1", "graph_enabled": True} + + async def variation(key: str, ctx: dict, default: Any) -> Any: + if key == "graph-key": + return ( + {"root": "only-node", "edges": {}} if state["graph_enabled"] else {} + ) + return { + "model": {"name": "gpt-4"}, + "provider": {"name": "TestProvider"}, + "instructions": state["instructions"], + "_ldMeta": {"enabled": True, "variationKey": "v1", "version": 1}, + } + + mock_ld_client.variation = AsyncMock(side_effect=variation) + seen: list[str] = [] + + async def fn( + config, user_input, tool_handlers, variables, history=None + ) -> dict: # type: ignore[no-untyped-def] + seen.append(config["instructions"]) + return {"output": "ok", "usage": {"input_tokens": 1, "output_tokens": 1}} + + g = graph( + "graph-key", + handlers=[ + ProviderHandler(fn=fn, provides_for=("TestProvider", "messages")) + ], # type: ignore[arg-type] + ) + await g.invoke("first", CONTEXT) + state["instructions"] = "v2" + await g.invoke("second", CONTEXT) + assert seen == ["v1", "v2"] + + state["graph_enabled"] = False + with pytest.raises(ValueError, match="disabled"): + await g.invoke("third", CONTEXT) + with pytest.raises(ValueError, match="disabled"): + async for _ in g.stream("fourth", CONTEXT): + pass + + async def test_failed_resolution_is_not_reused( + self, mock_ld_client: MagicMock + ) -> None: + """A transient variation failure on one call does not stick to the next.""" + original = mock_ld_client.variation.side_effect + failures = {"left": 1} + + async def flaky(key: str, ctx: dict, default: Any) -> Any: + if key != "graph-key" and failures["left"]: + failures["left"] -= 1 + raise RuntimeError("transient") + return await original(key, ctx, default) + + mock_ld_client.variation = AsyncMock(side_effect=flaky) + g = graph("graph-key", handlers=[_make_handler()]) + with pytest.raises(ValueError, match="disabled"): + await g.invoke("first", CONTEXT) + result = await g.invoke("second", CONTEXT) + assert result.response == "ok" async def test_handoff_success_tracked_on_each_handoff( self, mock_ld_client: MagicMock @@ -464,35 +530,6 @@ async def test_disabled_graph_reverse_traverse_key_is_snake_case( assert hasattr(gd, "reverse_traverse") assert not hasattr(gd, "reverseTraverse") - async def test_cache_is_bounded(self, mock_ld_client: MagicMock) -> None: - """GraphInstance._cache must not grow beyond MAX_GRAPH_CACHE_SIZE. - See TESTING.md §3.11.""" - from launchdarkly_ai_server.graph import MAX_GRAPH_CACHE_SIZE - - g = graph("graph-key", handlers=[_make_handler()]) - # Invoke with more distinct contexts than the max - limit = MAX_GRAPH_CACHE_SIZE + 10 - for i in range(limit): - await g.invoke("hi", {"kind": "user", "key": f"u{i}"}) - assert len(g._cache) <= MAX_GRAPH_CACHE_SIZE - - async def test_equal_content_contexts_share_cache_entry( - self, mock_ld_client: MagicMock - ) -> None: - """Two distinct dict objects with identical JSON content must share one - cache entry. The cache key must be json.dumps(context, sort_keys=True), - not id(context). See TESTING.md §3.11.""" - g = graph("graph-key", handlers=[_make_handler()]) - ctx_a = {"kind": "user", "key": "u1"} - ctx_b = {"kind": "user", "key": "u1"} # distinct object, same content - assert ctx_a is not ctx_b, "test requires two different dict objects" - await g.invoke("first call", ctx_a) - await g.invoke("second call", ctx_b) - assert len(g._cache) == 1, ( - "Equal-content contexts must share a single cache entry — " - "the cache key must be json.dumps(context), not id(context)." - ) - async def test_error_logging_when_node_variation_fails( self, mock_ld_client: MagicMock ) -> None: